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