Merge pull request #56 from smcazares/672-text-image-pair

672 text image pair

Author: smcazares

components/common/src/common/models/llm_query.py

@@ -201,12 +201,17 @@ class QueryReference(BaseModel):
   document_id = TextField(required=True)  # All modalities
   document_url = TextField(required=True)  # All modalities
   modality = TextField(required=True)  # All modalities: text image video audio
-  chunk_id = TextField(required=False)  # All modalities
-  chunk_url = TextField(required=False)  # Image or video or audio only
-  page = NumberField(required=False)  # Text or image only
-  document_text = TextField(required=False)  # Text only
-  timestamp_start = NumberField(required=False)  # Video or audio only
-  timestamp_stop = NumberField(required=False)  # Video or audio only
+  chunk_id = TextField(required=True)  # All modalities
+  chunk_url = TextField(
+    required=False, default=None)  # Image or video or audio only
+  page = NumberField(required=False, default=None)  # Text or image only
+  document_text = TextField(required=False, default=None)  # Text only
+  timestamp_start = NumberField(
+    required=False, default=None)  # Video or audio only
+  timestamp_stop = NumberField(
+    required=False, default=None)  # Video or audio only
+  linked_ids = ListField(
+    IDField(), required=False, default=None)  # All modalities
 
   def __repr__(self) -> str:
     """
@@ -218,24 +223,23 @@ def __repr__(self) -> str:
       document_text_snippet = self.document_text[:min(100,
                                                       document_text_num_chars)]
       chunk_url = None
-      page = None
     else:
       document_text_num_tokens = None
       document_text_num_chars = None
       document_text_snippet = None
       chunk_url = self.chunk_url
-      page = self.page
     return (
       f"Query_Ref(query_engine_name={self.query_engine}, "
       f"document_id={self.document_id}, "
       f"document_url={self.document_url}, "
       f"chunk_id={self.chunk_id}, "
       f"chunk_url={chunk_url}, "
       f"modality={self.modality}, "
-      f"page={page}, "
+      f"page={self.page}, "
       f"chunk_num_tokens={document_text_num_tokens}, "
       f"chunk_num_chars={document_text_num_chars}, "
-      f"chunk_text={document_text_snippet})"
+      f"chunk_text={document_text_snippet}, "
+      f"linked_ids={self.linked_ids})"
     )
 
   class Meta:
@@ -360,14 +364,18 @@ class QueryDocumentChunk(BaseModel):
   query_document_id = TextField(required=True)  # All modalities
   index = NumberField(required=True)  # All modalities
   modality = TextField(required=True)  # All modalities: text image video audio
-  page = NumberField(required=False)  # Text or image only
-  chunk_url = TextField(required=False)  # Image or video or audio only
-  text = TextField(required=False)  # Text only
-  clean_text = TextField(required=False)  # Text only (optional)
-  sentences = ListField(required=False)  # Text only (optional)
-  timestamp_start = NumberField(required=False)  # Video or audio only
-  timestamp_stop = NumberField(required=False)  # Video or audio only
-  linked_ids = ListField(required=False)  # All modalities
+  page = NumberField(required=False, default=None)  # Text or image only
+  chunk_url = TextField(
+    required=False, default=None)  # Image or video or audio only
+  text = TextField(required=False, default=None)  # Text only
+  clean_text = TextField(required=False, default=None)  # Text only (optional)
+  sentences = ListField(required=False, default=None)  # Text only (optional)
+  timestamp_start = NumberField(
+    required=False, default=None)  # Video or audio only
+  timestamp_stop = NumberField(
+    required=False, default=None)  # Video or audio only
+  linked_ids = ListField(
+    IDField(), required=False, default=None)  # All modalities
 
   class Meta:
     ignore_none_field = False

components/frontend_streamlit/src/pages/4_Query.py

@@ -95,20 +95,26 @@ def chat_content():
               _, chunk_type = splitext(chunk_url)
               chunk_url = chunk_url.replace("gs://",
                   "https://storage.googleapis.com/", 1)
-
             document_url = reference["document_url"]
+            page = reference["page"]
+            if page:
+              # References from multimodal query engines have page numbers
+              reference_header = (f"\nReference {query_index}:"
+                                  f" {document_url}, Page {page+1}")
+            else:
+              # References from text-only query engines do not have page numbers
+              reference_header = f"\nReference {query_index}: {document_url}"
             if modality == "text":
               document_text = reference["document_text"]
               st.text_area(
-                f"\nReference {query_index}: {document_url}",
+                reference_header,
                 document_text,
                 key=f"ref_{query_index}")
             elif modality == "image" and chunk_type in [".pdf",
                  ".png", ".jpg", ".jpeg", ".gif", ".bmp"]:
               # .tif/.tiff not available, all other file types are untested
-              page = reference["page"]
               st.write(
-                f"\nReference {query_index}: {document_url}, Page {page+1}",
+                reference_header,
                 key=f"ref_{query_index}")
               st.image(chunk_url)
             else:

components/llm_service/src/services/query/query_service.py

@@ -422,6 +422,20 @@ async def query_search(q_engine: QueryEngine,
     query_reference.save()
     query_references.append(query_reference)
 
+    # Also create a query_reference for other modalities of the same chunk
+    linked_ids = query_reference.linked_ids
+    if linked_ids:
+      for linked_id in linked_ids:
+        query_doc_chunk_friend = QueryDocumentChunk.find_by_id(linked_id)
+        query_reference_friend = make_query_reference(
+          q_engine=q_engine,
+          query_doc=query_doc,
+          doc_chunk=query_doc_chunk_friend,
+          query_embeddings=query_embeddings,
+          rank_sentences=rank_sentences)
+        query_reference_friend.save()
+        query_references.append(query_reference_friend)
+
   Logger.info(f"Retrieved {len(query_references)} "
                f"references={query_references}")
 
@@ -504,6 +518,7 @@ def make_query_reference(q_engine: QueryEngine,
   query_reference_dict["document_url"]=query_doc.doc_url
   query_reference_dict["modality"]=modality
   query_reference_dict["chunk_id"]=doc_chunk.id
+  query_reference_dict["linked_ids"]=doc_chunk.linked_ids
   # For text chunk only
   if modality=="text":
     query_reference_dict["page"]=doc_chunk.page
@@ -1071,7 +1086,10 @@ async def process_documents(doc_url: str, qe_vector_store: VectorStore,
           for index in linked_indexes:
             query_doc_chunk = \
               QueryDocumentChunk.find_by_index(q_engine.id, index)
-            query_doc_chunk.linked_ids = linked_ids
+            friend_ids = [friend_id for friend_id in linked_ids
+                          if friend_id != query_doc_chunk.id]
+            query_doc_chunk.linked_ids = friend_ids
+            query_doc_chunk.save()
 
         else:
           # Use text-only pipeline