image retrieval

z_generate.py

@@ -10,7 +10,7 @@ class ServerlessInference:
 
     def test(self, query:str) -> str:
         '''Responds to query using llm'''
-        messages:str = [
+        messages:list = [
             {
                 "role": "user",
                 "content": query
@@ -33,7 +33,7 @@ class ServerlessInference:
         context += "".join([f"Document {str(i)}:::\n" + doc for i, doc in enumerate(retrieved_docs_text)])
 
         # Augmented Generation
-        messages:str  = [
+        messages:list  = [
             {
                 "role": "system",
                 "content": """Using the information contained in the context,
@@ -42,9 +42,7 @@ give a comprehensive answer to the question.
 
 Respond only to the question asked, response should be concise and relevant to the question.
 
-Provide the number of the source document when relevant.
-
-If the answer cannot be deduced from the context, do not give an answer.""",
+If the answer cannot be deduced from the context, do not give an answer. Instead say `Theres lack of information in document source.`""",
 
             },
 
@@ -70,4 +68,49 @@ Question: {question}""".format(context=context, question=query),
         )
 
         response_text = completion.choices[0].message.content
-        return response_text
+
+        # Image retrieval 
+        retrieved_image = self.vs_images.similarity_search(query=query, k=5)
+        retrieved_docs_text = [doc.page_content for doc in retrieved_image]  # We only need the text of the documents
+        context = "\nExtracted Images:\n"
+        context += "".join([f"Document {str(i)}:::\n" + doc for i, doc in enumerate(retrieved_docs_text)])
+
+
+        messages:list  = [
+            {
+                "role": "system",
+                "content": """Using the information contained in the context about the images stored in the database,
+
+give a list of identifiers of the image that best represent the kind of information seeked by the question.
+
+Respond only to the question asked. Provide only number(s) of the source images relevant to the question. 
+
+If the image is relevant to the question then output format should be a list [1, 3, 0] 
+
+otherwise reply with []  (empty list)""",
+
+            },
+
+            {
+                "role": "user",
+                "content": """Context:
+
+{context}
+
+---
+
+Now here is the question you need to answer.
+
+Question: {question}""".format(context=context, question=query),
+
+            },
+        ]
+
+        completion = self.client.chat.completions.create(
+            model=self.model,
+            messages=messages, 
+            max_tokens=500
+        )
+
+        images_list = completion.choices[0].message.content
+        return response_text + str(images_list)