Pash1986 commited on
Commit
de24151
1 Parent(s): ee38225

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +147 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,147 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pymongo import MongoClient
2
+ import os
3
+ import time
4
+ import gradio as gr
5
+ import requests
6
+ import traceback
7
+ import google.generativeai as genai
8
+
9
+
10
+
11
+
12
+ genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
13
+ try:
14
+ # Initialize MongoDB python client
15
+ MONGODB_URI = os.getenv("MONGODB_ATLAS_URI")
16
+ client = MongoClient(MONGODB_URI, appname="devrel.content.python")
17
+
18
+ DB_NAME = "google-ai"
19
+ COLLECTION_NAME = "embedded_docs"
20
+ ATLAS_VECTOR_SEARCH_INDEX_NAME = "vector_index"
21
+ collection = client[DB_NAME][COLLECTION_NAME]
22
+
23
+ ### Insert data about 5 individual employees
24
+ collection.delete_many({})
25
+ collection.insert_many([
26
+ {
27
+ '_id' : '54633',
28
+ 'content' : 'Employee number 54633, name John Doe, department Sales, location New York, salary 100000'
29
+ },
30
+ {
31
+ '_id' : '54634',
32
+ 'content' : 'Employee number 54634, name Jane Doe, department Marketing, location Los Angeles, salary 120000',
33
+
34
+ },
35
+ {
36
+ '_id' : '54635',
37
+ 'content' : 'Employee number 54635, name John Smith, department Engineering, location San Francisco, salary 150000'
38
+ },
39
+ {
40
+ '_id' : '54636',
41
+ 'content' : 'Employee number 54636, name Jane Smith, department Finance, location Chicago, salary 130000'
42
+ },
43
+ {
44
+ '_id' : '54637',
45
+ 'content' : 'Employee number 54637, name John Johnson, department HR, location Miami, salary 110000'
46
+ },
47
+ {
48
+ '_id' : '54638',
49
+ 'content' : 'Employee number 54638, name Jane Johnson, department Operations, location Seattle, salary 140000'
50
+ }
51
+ ])
52
+
53
+
54
+ # Exception handling to catch and display errors during the pipeline execution.
55
+ except Exception as erorr_message:
56
+ print("An error occurred: \n" + erorr_message)
57
+
58
+ gemini_pro = genai.GenerativeModel('gemini-pro')
59
+
60
+ def embed_text(text):
61
+ result = genai.embed_content(
62
+ model="models/embedding-001",
63
+ content=text,
64
+ task_type="retrieval_document",
65
+ title="Embedding of single string")
66
+
67
+ return result['embedding']
68
+
69
+ def get_rag_output(context, question):
70
+
71
+ template = f""" You are an hr assistant, answer in detail. Answer the question based only on the following context:
72
+ ```
73
+ {context}
74
+ ```
75
+ Question: {question}
76
+ """
77
+ response = gemini_pro.generate_content([template], stream=False)
78
+ return response.text
79
+
80
+ def mongodb_vector_query(message):
81
+ docs = collection.aggregate([
82
+ {
83
+ '$vectorSearch' : {
84
+ 'index' : 'vector_index',
85
+ 'queryVector' : embed_text(message),
86
+ 'path' : 'embedding',
87
+ 'numCandidates' : 10,
88
+ 'limit' : 5
89
+ }
90
+ },
91
+ {
92
+ '$project': {
93
+ 'embedding': 0
94
+ }
95
+ }
96
+ ])
97
+
98
+ return list(docs)
99
+
100
+
101
+
102
+
103
+ def get_rag(message, history):
104
+
105
+ try:
106
+ context = mongodb_vector_query(message)
107
+ result = get_rag_output(context, message)
108
+
109
+ # print(result)
110
+ print_llm_text = result
111
+ for i in range(len(print_llm_text)):
112
+ time.sleep(0.03)
113
+ yield print_llm_text[: i+1]
114
+
115
+
116
+ except Exception as e:
117
+ error_message = traceback.format_exc()
118
+ print("An error occurred: \n" + error_message)
119
+ yield error_message
120
+
121
+
122
+
123
+
124
+ def fetch_url_data(url):
125
+ try:
126
+ response = requests.get(url)
127
+ response.raise_for_status() # Raises an HTTPError if the HTTP request returned an unsuccessful status code
128
+ return response.text
129
+ except requests.RequestException as e:
130
+ return f"Error: {e}"
131
+
132
+ # Setup Gradio interface
133
+ with gr.Blocks() as demo:
134
+ with gr.Tab("Demo"):
135
+
136
+ ## value=[(None, "Hi, I'm a MongoDB and Heystack based question and answer bot 🤖, I can help you answer on the knowledge base above…")]
137
+ gr.ChatInterface(get_rag,examples=["List all employees", "Where does jane work?", "Who has the highest salary? List it"], title="Atlas Vector Search Chat",description="This small chat uses a similarity search to find relevant plots as listed above, it uses MongoDB Atlas and Google Gemini.",submit_btn="Search").queue()
138
+
139
+ with gr.Tab("Code"):
140
+ gr.Code(label="Code", language="python", value=fetch_url_data('https://huggingface.co/spaces/MongoDB/Haystack-MongoDB-Integration-Chat/raw/main/app.py'))
141
+
142
+
143
+
144
+
145
+
146
+ if __name__ == "__main__":
147
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ google-generativeai
2
+ pymongo
3
+ requests