bstraehle commited on
Commit
7f1903e
·
verified ·
1 Parent(s): 6608ce6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -103
app.py CHANGED
@@ -4,7 +4,7 @@ import logging, os, sys, threading
4
 
5
  from datasets import load_dataset
6
  from dotenv import load_dotenv, find_dotenv
7
- from utils import process_records, connect_to_database, setup_vector_search_index
8
 
9
  from pydantic import BaseModel
10
  from typing import Optional
@@ -23,108 +23,6 @@ RAG_ADVANCED = "Advanced RAG"
23
 
24
  logging.basicConfig(stream = sys.stdout, level = logging.INFO)
25
  logging.getLogger().addHandler(logging.StreamHandler(stream = sys.stdout))
26
-
27
- def vector_search(user_query, db, collection, additional_stages=[], vector_index="vector_index_text"):
28
- """
29
- Perform a vector search in the MongoDB collection based on the user query.
30
-
31
- Args:
32
- user_query (str): The user's query string.
33
- db (MongoClient.database): The database object.
34
- collection (MongoCollection): The MongoDB collection to search.
35
- additional_stages (list): Additional aggregation stages to include in the pipeline.
36
-
37
- Returns:
38
- list: A list of matching documents.
39
- """
40
-
41
- # Generate embedding for the user query
42
- query_embedding = custom_utils.get_embedding(user_query)
43
-
44
- if query_embedding is None:
45
- return "Invalid query or embedding generation failed."
46
-
47
- # Define the vector search stage
48
- vector_search_stage = {
49
- "$vectorSearch": {
50
- "index": vector_index, # specifies the index to use for the search
51
- "queryVector": query_embedding, # the vector representing the query
52
- "path": "text_embeddings", # field in the documents containing the vectors to search against
53
- "numCandidates": 150, # number of candidate matches to consider
54
- "limit": 20, # return top 20 matches
55
- }
56
- }
57
-
58
- # Define the aggregate pipeline with the vector search stage and additional stages
59
- pipeline = [vector_search_stage] + additional_stages
60
-
61
- # Execute the search
62
- results = collection.aggregate(pipeline)
63
-
64
- explain_query_execution = db.command( # sends a database command directly to the MongoDB server
65
- 'explain', { # return information about how MongoDB executes a query or command without actually running it
66
- 'aggregate': collection.name, # specifies the name of the collection on which the aggregation is performed
67
- 'pipeline': pipeline, # the aggregation pipeline to analyze
68
- 'cursor': {} # indicates that default cursor behavior should be used
69
- },
70
- verbosity='executionStats') # detailed statistics about the execution of each stage of the aggregation pipeline
71
-
72
- vector_search_explain = explain_query_execution['stages'][0]['$vectorSearch']
73
- millis_elapsed = vector_search_explain['explain']['collectStats']['millisElapsed']
74
-
75
- print(f"Total time for the execution to complete on the database server: {millis_elapsed} milliseconds")
76
-
77
- return list(results)
78
-
79
- class SearchResultItem(BaseModel):
80
- name: str
81
- accommodates: Optional[int] = None
82
- bedrooms: Optional[int] = None
83
- address: custom_utils.Address
84
- space: str = None
85
-
86
- def handle_user_query(query, db, collection, stages=[], vector_index="vector_index_text"):
87
- # Assuming vector_search returns a list of dictionaries with keys 'title' and 'plot'
88
- get_knowledge = vector_search(query, db, collection, stages, vector_index)
89
-
90
- # Check if there are any results
91
- if not get_knowledge:
92
- return "No results found.", "No source information available."
93
-
94
- # Convert search results into a list of SearchResultItem models
95
- search_results_models = [
96
- SearchResultItem(**result)
97
- for result in get_knowledge
98
- ]
99
-
100
- # Convert search results into a DataFrame for better rendering in Jupyter
101
- search_results_df = pd.DataFrame([item.dict() for item in search_results_models])
102
-
103
- # Generate system response using OpenAI's completion
104
- completion = custom_utils.openai.chat.completions.create(
105
- model="gpt-3.5-turbo",
106
- messages=[
107
- {
108
- "role": "system",
109
- "content": "You are a airbnb listing recommendation system."},
110
- {
111
- "role": "user",
112
- "content": f"Answer this user query: {query} with the following context:\n{search_results_df}"
113
- }
114
- ]
115
- )
116
-
117
- system_response = completion.choices[0].message.content
118
-
119
- # Print User Question, System Response, and Source Information
120
- print(f"- User Question:\n{query}\n")
121
- print(f"- System Response:\n{system_response}\n")
122
-
123
- # Display the DataFrame as an HTML table
124
- display(HTML(search_results_df.to_html()))
125
-
126
- # Return structured response and source info as a string
127
- return system_response
128
 
129
  def invoke(openai_api_key, prompt, rag_option):
130
  if not openai_api_key:
 
4
 
5
  from datasets import load_dataset
6
  from dotenv import load_dotenv, find_dotenv
7
+ from custom_utils import process_records, connect_to_database, setup_vector_search_index
8
 
9
  from pydantic import BaseModel
10
  from typing import Optional
 
23
 
24
  logging.basicConfig(stream = sys.stdout, level = logging.INFO)
25
  logging.getLogger().addHandler(logging.StreamHandler(stream = sys.stdout))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26
 
27
  def invoke(openai_api_key, prompt, rag_option):
28
  if not openai_api_key: