Update custom_utils.py

custom_utils.py

@@ -29,20 +29,16 @@ def rag_ingestion(collection):
     return "Manually create a vector search index (in free tier, this feature is not available via SDK)"
 
 def rag_retrieval(openai_api_key, prompt, db, collection, stages=[], vector_index="vector_index"):
-    # Assuming vector_search returns a list of dictionaries with keys 'title' and 'plot'
     get_knowledge = vector_search(openai_api_key, prompt, db, collection, stages, vector_index)
 
-    # Check if there are any results
     if not get_knowledge:
         return "No results found.", "No source information available."
 
-    # Convert search results into a list of SearchResultItem models
     search_results_models = [
         SearchResultItem(**result)
         for result in get_knowledge
     ]
 
-    # Convert search results into a DataFrame for better rendering in Jupyter
     search_results_df = pd.DataFrame([item.dict() for item in search_results_models])
 
     print("###")
@@ -54,7 +50,6 @@ def rag_retrieval(openai_api_key, prompt, db, collection, stages=[], vector_inde
 def rag_inference(openai_api_key, prompt, search_results):
     openai.api_key = openai_api_key
     
-    # Generate system response using OpenAI's completion
     content = f"Answer this user question: {prompt} with the following context:\n{search_results}"
     
     completion = openai.chat.completions.create(
@@ -70,13 +65,7 @@ def rag_inference(openai_api_key, prompt, search_results):
         ]
     )
 
-    completion_result = completion.choices[0].message.content
-
-    print("###")
-    print(completion_result)
-    print("###")
-
-    return completion_result
+    return completion.choices[0].message.content
 
 def process_records(data_frame):
     records = data_frame.to_dict(orient="records")
@@ -101,33 +90,19 @@ def process_records(data_frame):
         return []
 
 def vector_search(openai_api_key, user_query, db, collection, additional_stages=[], vector_index="vector_index"):
-    """
-    Perform a vector search in the MongoDB collection based on the user query.
-
-    Args:
-    user_query (str): The user's query string.
-    db (MongoClient.database): The database object.
-    collection (MongoCollection): The MongoDB collection to search.
-    additional_stages (list): Additional aggregation stages to include in the pipeline.
-
-    Returns:
-    list: A list of matching documents.
-    """
-
-    # Generate embedding for the user query
-    query_embedding = get_embedding(openai_api_key, user_query)
+    query_embedding = get_text_embedding(openai_api_key, user_query)
+    query_embedding2 = get_image_embedding(openai_api_key, user_query)
 
     if query_embedding is None:
         return "Invalid query or embedding generation failed."
 
-    # Define the vector search stage
     vector_search_stage = {
         "$vectorSearch": {
-            "index": vector_index, # specifies the index to use for the search
-            "queryVector": query_embedding, # the vector representing the query
-            "path": "text_embeddings", # field in the documents containing the vectors to search against
-            "numCandidates": 150, # number of candidate matches to consider
-            "limit": 20, # return top 20 matches
+            "index": vector_index,
+            "queryVector": query_embedding,
+            "path": "text_embeddings",
+            "numCandidates": 150,
+            "limit": 20,
             "filter": {
                 "$and": [
                     {"accommodates": {"$eq": 2}}, 
@@ -137,36 +112,25 @@ def vector_search(openai_api_key, user_query, db, collection, additional_stages=
         }
     }
     
-    # Define the aggregate pipeline with the vector search stage and additional stages
     pipeline = [vector_search_stage] + additional_stages
 
-    # Execute the search
     results = collection.aggregate(pipeline)
 
-    explain_query_execution = db.command( # sends a database command directly to the MongoDB server
-        'explain', { # return information about how MongoDB executes a query or command without actually running it
-            'aggregate': collection.name, # specifies the name of the collection on which the aggregation is performed
-            'pipeline': pipeline, # the aggregation pipeline to analyze
-            'cursor': {} # indicates that default cursor behavior should be used
+    explain_query_execution = db.command(
+        "explain", {
+            "aggregate": collection.name,
+            "pipeline": pipeline,
+            "cursor": {}
         }, 
-        verbosity='executionStats') # detailed statistics about the execution of each stage of the aggregation pipeline
+        verbosity='executionStats')
 
-    vector_search_explain = explain_query_execution['stages'][0]['$vectorSearch']
-
-    millis_elapsed = vector_search_explain['explain']['collectStats']['millisElapsed']
-    print("###")
-    print(vector_search_explain)
-    print("###")
-    print(vector_search_explain['explain'])
-    print("###")
-    print(f"Total time for the execution to complete on the database server: {millis_elapsed} milliseconds")
+    vector_search_explain = explain_query_execution["stages"][0]["$vectorSearch"]
+    millis_elapsed = vector_search_explain["explain"]["collectStats"]["millisElapsed"]
+    print(f"Query execution time: {millis_elapsed} milliseconds")
 
     return list(results)
 
-def get_embedding(openai_api_key, text):
-    """Generate an embedding for the given text using OpenAI's API."""
-
-    # Check for valid input
+def get_text_embedding(openai_api_key, text):
     if not text or not isinstance(text, str):
         return None
 
@@ -177,6 +141,21 @@ def get_embedding(openai_api_key, text):
             input=text,
             model="text-embedding-3-small", dimensions=1536).data[0].embedding
         return embedding
+    except Exception as e:
+        print(f"Error in get_embedding: {e}")
+        return None
+
+def get_image_embedding(openai_api_key, text):
+    if not text or not isinstance(text, str):
+        return None
+
+    openai.api_key = openai_api_key
+        
+    try:
+        embedding = openai.embeddings.create(
+            input=text,
+            model="openai/clip-vit-base-patch32", dimensions=512.data[0].embedding
+        return embedding
     except Exception as e:
         print(f"Error in get_embedding: {e}")
         return None