Merge pull request #54 from GPS-Solutions/resolve_merge_conflicts_from_main_to_dev

Merge aica:main to aica:dev

Author: smcazares

components/llm_service/src/routes/chat.py

@@ -273,17 +273,15 @@ async def create_user_chat(
   if prompt is None or prompt == "":
     return BadRequest("Missing or invalid payload parameters")
 
-  # process chat file: upload to GCS and determine mime type
-  chat_file_type = None
+  # process chat file(s): upload to GCS and determine mime type
   chat_file_bytes = None
-  chat_file_urls = None
+  chat_files = None
   if chat_file is not None or chat_file_url is not None:
-    chat_file_urls, chat_file_type = \
-        await process_chat_file(chat_file, chat_file_url)
+    chat_files = await process_chat_file(chat_file, chat_file_url)
 
   # only read chat file bytes if for some reason we can't
-  # upload the file to GCS
-  if not chat_file_urls and chat_file is not None:
+  # upload the file(s) to GCS
+  if not chat_files and chat_file is not None:
     await chat_file.seek(0)
     chat_file_bytes = await chat_file.read()
 
@@ -293,9 +291,8 @@ async def create_user_chat(
     # generate text from prompt
     response = await llm_chat(prompt,
                               llm_type,
-                              chat_file_types=[chat_file_type],
-                              chat_file_bytes=chat_file_bytes,
-                              chat_file_urls=chat_file_urls)
+                              chat_files=chat_files,
+                              chat_file_bytes=chat_file_bytes)
 
     # create new chat for user
     user_chat = UserChat(user_id=user.user_id, llm_type=llm_type,

components/llm_service/src/services/llm_generate.py

@@ -38,6 +38,7 @@
                     KEY_MODEL_PARAMS, KEY_MODEL_CONTEXT_LENGTH,
                     DEFAULT_LLM_TYPE, DEFAULT_MULTIMODAL_LLM_TYPE)
 from services.langchain_service import langchain_llm_generate
+from services.query.data_source import DataSourceFile
 from utils.errors import ContextWindowExceededException
 
 Logger = Logger.get_logger(__file__)
@@ -107,17 +108,16 @@ async def llm_generate(prompt: str, llm_type: str) -> str:
   except Exception as e:
     raise InternalServerError(str(e)) from e
 
-async def llm_generate_multimodal(prompt: str, llm_type: str, user_file_types: List[str],
+async def llm_generate_multimodal(prompt: str, llm_type: str,
                              user_file_bytes: bytes = None,
-                             user_file_urls: List[str] = None) -> str:
+                             user_files: List[DataSourceFile] = None) -> str:
   """
   Generate text with an LLM given a file and a prompt.
   Args:
     prompt: the text prompt to pass to the LLM
-    user_file_bytes: bytes of the file provided by the user
-    user_file_urls: list of URLs to include in context
-    user_file_types: list of mime times for files to include in context
     llm_type: the type of LLM to use (default to gemini)
+    user_file_bytes: bytes of the file provided by the user
+    user_files: list of DataSourceFile objects for file meta data
   Returns:
     the text response: str
   """
@@ -146,7 +146,7 @@ async def llm_generate_multimodal(prompt: str, llm_type: str, user_file_types: L
             f"Vertex model {llm_type} needs to be multimodal")
       response = await google_llm_predict(prompt, is_chat, is_multimodal,
                             google_llm, None, user_file_bytes,
-                            user_file_urls, user_file_types)
+                            user_files)
     else:
       raise ResourceNotFoundException(f"Cannot find llm type '{llm_type}'")
 
@@ -160,8 +160,7 @@ async def llm_generate_multimodal(prompt: str, llm_type: str, user_file_types: L
 async def llm_chat(prompt: str, llm_type: str,
                    user_chat: Optional[UserChat] = None,
                    user_query: Optional[UserQuery] = None,
-                   chat_file_types: Optional[List[str]] = None,
-                   chat_file_urls: Optional[List[str]] = None,
+                   chat_files: Optional[List[DataSourceFile]] = None,
                    chat_file_bytes: Optional[bytes] = None) -> str:
   """
   Send a prompt to a chat model and return string response.
@@ -173,9 +172,9 @@ async def llm_chat(prompt: str, llm_type: str,
     llm_type: the type of LLM to use
     user_chat (optional): a user chat to use for context
     user_query (optional): a user query to use for context
-    chat_file_bytes (bytes): bytes of file to include in chat context
-    chat_file_urls (List[str]): urls of files to include in chat context
-    chat_file_types (List[str]): mime types of files to include in chat context
+    chat_files (optional) (List[DataSourceFile]): files to include in chat context
+    chat_file_bytes (optional) (bytes): bytes of file to include in chat context
+
   Returns:
     the text response: str
   """
@@ -185,22 +184,19 @@ async def llm_chat(prompt: str, llm_type: str,
               f" user_chat=[{user_chat}]"
               f" user_query=[{user_query}]"
               f" chat_file_bytes=[{chat_file_bytes_log}]"
-              f" chat_file_urls=[{chat_file_urls}]"
-              f" chat_file_type=[{chat_file_types}]")
+              f" chat_files=[{chat_files}]")
 
   if llm_type not in get_model_config().get_chat_llm_types():
     raise ResourceNotFoundException(f"Cannot find chat llm type '{llm_type}'")
 
   # validate chat file params
   is_multimodal = False
-  if chat_file_bytes is not None or chat_file_urls:
-    if chat_file_bytes is not None and chat_file_urls:
+  if chat_file_bytes is not None or chat_files:
+    if chat_file_bytes is not None and chat_files:
       raise InternalServerError(
-          "Must set only one of chat_file_bytes/chat_file_urls")
+          "Must set only one of chat_file_bytes/chat_files")
     if llm_type not in get_provider_models(PROVIDER_VERTEX):
       raise InternalServerError("Chat files only supported for Vertex")
-    if chat_file_types is None:
-      raise InternalServerError("Mime type must be passed for chat file")
     is_multimodal = True
 
   try:
@@ -244,7 +240,7 @@ async def llm_chat(prompt: str, llm_type: str,
       response = await google_llm_predict(prompt, is_chat, is_multimodal,
                                           google_llm, user_chat,
                                           chat_file_bytes,
-                                          chat_file_urls, chat_file_types)
+                                          chat_files)
     elif llm_type in get_provider_models(PROVIDER_LANGCHAIN):
       response = await langchain_llm_generate(prompt, llm_type, user_chat)
     return response
@@ -496,9 +492,8 @@ async def model_garden_predict(prompt: str,
 
 async def google_llm_predict(prompt: str, is_chat: bool, is_multimodal: bool,
                 google_llm: str, user_chat=None,
-                user_file_bytes: Optional[bytes]=None,
-                user_file_urls: Optional[List[str]]=None,
-                user_file_types: Optional[List[str]]=None) -> str:
+                user_file_bytes: bytes=None,
+                user_files: List[DataSourceFile]=None) -> str:
   """
   Generate text with a Google multimodal LLM given a prompt.
   Args:
@@ -508,8 +503,7 @@ async def google_llm_predict(prompt: str, is_chat: bool, is_multimodal: bool,
     google_llm: name of the vertex llm model
     user_chat: chat history
     user_file_bytes: the bytes of the file provided by the user
-    user_file_urls: list of urls of files provided by the user
-    user_file_types: list of mime types of the files provided by the user
+    user_files: list of DataSourceFiles for files provided by the user
   Returns:
     the text response.
   """
@@ -519,8 +513,7 @@ async def google_llm_predict(prompt: str, is_chat: bool, is_multimodal: bool,
               f" prompt=[{prompt}], is_chat=[{is_chat}],"
               f" is_multimodal=[{is_multimodal}], google_llm=[{google_llm}],"
               f" user_file_bytes=[{user_file_bytes_log}],"
-              f" user_file_urls=[{user_file_urls}],"
-              f" user_file_type=[{user_file_types}].")
+              f" user_files=[{user_files}]")
 
   # TODO: Consider images in chat
   prompt_list = []
@@ -563,21 +556,20 @@ async def google_llm_predict(prompt: str, is_chat: bool, is_multimodal: bool,
         chat_model = GenerativeModel(google_llm)
         if is_multimodal:
           user_file_parts = []
-          if user_file_bytes is not None:
+          if user_file_bytes is not None and user_files is not None:
             # user_file_bytes refers to a single image and so we index into
-            # user_file_types (a list) to get a single mime type
+            # user_files (a list) to get a single mime type
             user_file_parts = [Part.from_data(user_file_bytes,
-                                              mime_type=user_file_types[0])]
-          elif user_file_urls is not None:
-            # user_file_urls and user_file_types are same-length lists
-            # referring to one or more images
+                                              mime_type=user_files[0].mime_type)]
+          elif user_files is not None:
+            # user_files is a list referring to one or more images
             user_file_parts = [
-              Part.from_uri(user_file_url, mime_type=user_file_type)
-              for user_file_url, user_file_type in zip(user_file_urls, user_file_types)
+              Part.from_uri(user_file.gcs_path, mime_type=user_file.mime_type)
+              for user_file in user_files
             ]
           else:
             raise RuntimeError(
-                "if is_multimodal one of user_file_bytes or user_file_urls must be set")
+                "if is_multi user_files must be set")
           context_list = [*user_file_parts, context_prompt]
           Logger.info(f"context list {context_list}")
           generation_config = GenerationConfig(**parameters)

components/llm_service/src/services/llm_generate_test.py

@@ -37,6 +37,7 @@
                                                clean_firestore)
 from common.utils.logging_handler import Logger
 from schemas.schema_examples import (CHAT_EXAMPLE, USER_EXAMPLE)
+from services.query.data_source import DataSourceFile
 
 Logger = Logger.get_logger(__file__)
 
@@ -77,7 +78,6 @@
 }
 
 FAKE_FILE_NAME = "test.png"
-FAKE_FILE_TYPE = "image/png"
 FAKE_PROMPT = "test prompt"
 
 
@@ -147,7 +147,7 @@ async def test_llm_generate_google(clean_firestore):
 
 
 @pytest.mark.asyncio
-async def test_llm_generate_multimodal(clean_firestore):
+async def test_llm_generate_multi_file(clean_firestore):
   get_model_config().llm_model_providers = {
     PROVIDER_VERTEX: TEST_VERTEX_CONFIG
   }
@@ -159,18 +159,38 @@ async def test_llm_generate_multimodal(clean_firestore):
   os.remove(FAKE_FILE_NAME)
   fake_upload_file = UploadFile(file=fake_file, filename=FAKE_FILE_NAME)
   fake_file_bytes = await fake_upload_file.read()
-
+  fake_file_data = [DataSourceFile(mime_type="image/png")]
   with mock.patch(
   "vertexai.preview.generative_models.GenerativeModel.generate_content_async",
   return_value=FAKE_GOOGLE_RESPONSE):
     response = await llm_generate_multimodal(FAKE_PROMPT,
                                         VERTEX_LLM_TYPE_GEMINI_PRO_VISION,
-                                        [FAKE_FILE_TYPE],
-                                        fake_file_bytes)
+                                        fake_file_bytes,
+                                        fake_file_data)
   fake_file.close()
   assert response == FAKE_GENERATE_RESPONSE
 
 
+@pytest.mark.asyncio
+async def test_llm_generate_multi_url(clean_firestore):
+  get_model_config().llm_model_providers = {
+    PROVIDER_VERTEX: TEST_VERTEX_CONFIG
+  }
+  get_model_config().llm_models = TEST_VERTEX_CONFIG
+
+  fake_file_data = [DataSourceFile(mime_type="image/png",
+                                   gcs_path="gs://fake_bucket/file.png")]
+  fake_file_bytes = None
+  with mock.patch(
+  "vertexai.preview.generative_models.GenerativeModel.generate_content_async",
+  return_value=FAKE_GOOGLE_RESPONSE):
+    response = await llm_generate_multimodal(FAKE_PROMPT,
+                                        VERTEX_LLM_TYPE_GEMINI_PRO_VISION,
+                                        fake_file_bytes,
+                                        fake_file_data)
+  assert response == FAKE_GENERATE_RESPONSE
+
+
 @pytest.mark.asyncio
 async def test_llm_chat_google(clean_firestore, test_chat):
   get_model_config().llm_model_providers = {

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

@@ -54,12 +54,28 @@ def __init__(self,
                src_url:str=None,
                local_path:str=None,
                gcs_path:str=None,
-               doc_id:str=None):
+               doc_id:str=None,
+               mime_type:str=None):
     self.doc_name = doc_name
     self.src_url = src_url
     self.local_path = local_path
     self.gcs_path = gcs_path
     self.doc_id = doc_id
+    self.mime_type = mime_type
+
+  def __repr__(self) -> str:
+    """
+    Log-friendly string representation of a DataSourceFile
+    """
+    return (
+      f"DataSourceFile(doc_name={self.doc_name}, "
+      f"src_url={self.src_url}, "
+      f"local_path={self.local_path}, "
+      f"gcs_path={self.gcs_path}, "
+      f"doc_id={self.doc_id}, "
+      f"mime_type={self.mime_type})"
+    )
+
 
 class DataSource:
   """

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

@@ -49,7 +49,7 @@
                                          MatchingEngineVectorStore,
                                          PostgresVectorStore,
                                          NUM_MATCH_RESULTS)
-from services.query.data_source import DataSource
+from services.query.data_source import DataSource, DataSourceFile
 from services.query.web_datasource import WebDataSource
 from services.query.sharepoint_datasource import SharePointDataSource
 from services.query.vertex_search import (build_vertex_search,
@@ -186,23 +186,21 @@ async def query_generate(
 
   # generate list of URLs for additional context
   # (from non-text info in query_references)
-  context_urls = []
-  context_urls_mimetype = []
+  context_files = []
   for ref in query_references:
     if hasattr(ref, "modality") and ref.modality != "text":
       if hasattr(ref, "chunk_url"):
         ref_filename = ref.chunk_url
         ref_mimetype = validate_multimodal_file_type(file_name=ref_filename,
                                                      file_b64=None)
-        context_urls.append(ref_filename)
-        context_urls_mimetype.append(ref_mimetype)
-        # TODO: If ref is a video chunk, then update ref.chunk_url
-        # according to ref.timestamp_start and ref.timestamp_stop
+        context_files.append(DataSourceFile(gcs_path=ref_filename,
+                                            mime_type=ref_mimetype))
+        # TODO: If ref is a video chunk, then update new element of
+        # context_files according to ref.timestamp_start and ref.timestamp_stop
 
   # send prompt and additional context to model
   question_response = await llm_chat(question_prompt, llm_type,
-                                     chat_file_types=context_urls_mimetype,
-                                     chat_file_urls=context_urls)
+                                     chat_files=context_files)
 
   # update user query with response
   if user_query:

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

@@ -25,7 +25,7 @@
 import sys
 import tempfile
 from pathlib import Path
-from typing import List
+from typing import List, Union
 from scrapy import signals
 from scrapy.crawler import CrawlerProcess
 from scrapy.linkextractors import LinkExtractor
@@ -42,14 +42,19 @@
 
 Logger = Logger.get_logger(__file__)
 
-def save_content(filepath: str, file_name: str, content: str) -> None:
+def save_content(filepath: str, file_name: str,
+                 content: Union[str, bytes]) -> None:
   """
   Save content in a file in a local directory
   """
   Logger.info(f"Saving {file_name} to {filepath}")
   doc_filepath = os.path.join(filepath, file_name)
-  with open(doc_filepath, "w", encoding="utf-8") as f:
-    f.write(content)
+  if isinstance(content, bytes):
+    with open(doc_filepath, "wb") as f:
+      f.write(content)
+  else:
+    with open(doc_filepath, "w", encoding="utf-8") as f:
+      f.write(content)
   Logger.info(f"{len(content)} bytes written")
   return doc_filepath