Merge branch 'rag-langchain-chat-history' of https://github.com/Googl…

…eCloudPlatform/ai-on-gke into rag-langchain-chat-history
GoogleCloudPlatform · Sep 9, 2024 · e209b46 · e209b46
2 parents 88ee300 + 324abf7
commit e209b46
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 144 deletions.
diff --git a/applications/rag/frontend/container/Dockerfile b/applications/rag/frontend/container/Dockerfile
@@ -22,6 +22,6 @@ RUN pip install -r requirements.txt
 EXPOSE 8080
 
 ENV FLASK_APP=/workspace/frontend/main.py
-ENV PYTHONPATH=.
+ENV PYTHONPATH=/workspace/frontend/
 # Run the application with Gunicorn
 CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:8080", "main:app"]
diff --git a/applications/rag/frontend/container/application/utils/cloud_sql_utils.py b/applications/rag/frontend/container/application/utils/cloud_sql_utils.py
@@ -22,13 +22,18 @@
 import os
 import logging
 
+import google.cloud.logging as gcloud_logging
+
 from google.cloud.sql.connector import IPTypes
 
-from langchain_google_cloud_sql_pg import PostgresEngine, PostgresVectorStore
+from langchain_google_cloud_sql_pg import PostgresEngine
 
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
 )
+gcloud_logging_client = gcloud_logging.Client()
+gcloud_logging_client.setup_logging()
+
 
 ENVIRONMENT = os.environ.get("ENVIRONMENT")
 
@@ -38,7 +43,7 @@
 
 DB_NAME = os.environ.get("DB_NAME", "pgvector-database")
 VECTOR_EMBEDDINGS_TABLE_NAME = os.environ.get("EMBEDDINGS_TABLE_NAME", "")
-CHAT_HISTORY_TABLE_NAME = os.environ.get("CHAT_HISTORY_TABLE_NAME", "chat_history_store")
+CHAT_HISTORY_TABLE_NAME = os.environ.get("CHAT_HISTORY_TABLE_NAME", "message_store")
 
 VECTOR_DIMENSION = os.environ.get("VECTOR_DIMENSION", 384)
 
@@ -66,13 +71,17 @@ def create_sync_postgres_engine():
         ip_type=IPTypes.PUBLIC if ENVIRONMENT == "development" else IPTypes.PRIVATE,
     )
     try:
-        engine.init_chat_history_table(table_name=CHAT_HISTORY_TABLE_NAME)
         engine.init_vectorstore_table(
             VECTOR_EMBEDDINGS_TABLE_NAME,
             vector_size=VECTOR_DIMENSION,
             overwrite_existing=False,
         )
-    except Exception as e:
-        logging.info(f"Error: {e}")
+    except Exception as err:
+        logging.error(f"Error: {err}")
+
+    try:
+        engine.init_chat_history_table(table_name=CHAT_HISTORY_TABLE_NAME)
 
+    except Exception as err:
+        logging.error(f"Error: {err}")
     return engine
diff --git a/applications/rag/frontend/container/application/vector_storages/cloud_sql.py b/applications/rag/frontend/container/application/vector_storages/cloud_sql.py
@@ -65,8 +65,8 @@ def add_texts(
             splits = self.splitter.split_documents(texts)
             ids = [str(uuid.uuid4()) for _ in range(len(splits))]
             self.vector_store.add_documents(splits, ids)
-        except Exception as e:
-            logging.info(f"Error: {e}")
+        except Exception as err:
+            logging.error(f"Error: {err}")
             raise Exception(f"Error adding texts: {err}")
 
     def similarity_search(
@@ -79,4 +79,5 @@ def similarity_search(
             return docs
 
         except Exception as err:
+            logging.error(f"Something happened: {err}")
             raise Exception(f"Error on similarity search: {err}")
diff --git a/applications/rag/frontend/container/rag_langchain/rag_chain.py b/applications/rag/frontend/container/rag_langchain/rag_chain.py