wip

Author: ks6088ts

.env.template

@@ -69,6 +69,17 @@ SQL_DATABASE_URI=""
 # SQL_DATABASE_URI="sqlite:///template_langgraph.db" # SQLite
 # SQL_DATABASE_URI="postgresql://user:password@localhost:5432/db" # PostgreSQL
 
+# PostgreSQL Settings
+POSTGRESQL_USER="user"
+POSTGRESQL_PASSWORD="password"
+POSTGRESQL_HOST="localhost"
+POSTGRESQL_PORT="5432"
+POSTGRESQL_DATABASE="db"
+POSTGRESQL_TABLE_NAME="reports_kabuto_vectors"
+POSTGRESQL_ID_COLUMN="id"
+POSTGRESQL_CONTENT_COLUMN="user_message"
+POSTGRESQL_EMBEDDING_COLUMN="embedding"
+
 # Azure AI Search Settings
 AI_SEARCH_ENDPOINT="https://xxx.search.windows.net/"
 AI_SEARCH_KEY="xxx"

scripts/postgresql_operator.py

@@ -0,0 +1,105 @@
+import logging
+
+import typer
+from dotenv import load_dotenv
+
+from template_langgraph.internals.csv_loaders import CsvLoaderWrapper
+from template_langgraph.internals.pdf_loaders import PdfLoaderWrapper
+from template_langgraph.loggers import get_logger
+from template_langgraph.tools.postgresql_tool import PostgreSQLClientWrapper
+
+# Initialize the Typer application
+app = typer.Typer(
+    add_completion=False,
+    help="PostgreSQL operator CLI",
+)
+
+# Set up logging
+logger = get_logger(__name__)
+
+
+@app.command()
+def add_documents(
+    verbose: bool = typer.Option(
+        False,
+        "--verbose",
+        "-v",
+        help="Enable verbose output",
+    ),
+):
+    # Set up logging
+    if verbose:
+        logger.setLevel(logging.DEBUG)
+
+    # Load documents from PDF files
+    pdf_documents = PdfLoaderWrapper().load_pdf_docs()
+    logger.info(f"Loaded {len(pdf_documents)} documents from PDF.")
+
+    # Load documents from CSV files
+    csv_documents = CsvLoaderWrapper().load_csv_docs()
+    logger.info(f"Loaded {len(csv_documents)} documents from CSV.")
+
+    # Combine all documents
+    documents = pdf_documents + csv_documents
+    logger.info(f"Total documents to add: {len(documents)}")
+
+    # Add documents to PostgreSQL
+    postgresql_client = PostgreSQLClientWrapper()
+    ids = postgresql_client.add_documents(
+        documents=documents,
+    )
+    logger.info(f"Added {len(ids)} documents to PostgreSQL.")
+    for id in ids:
+        logger.debug(f"Added document ID: {id}")
+
+    # FIXME: deleting documents() is not working as expected
+    # assert postgresql_client.delete_documents(ids=ids), "Failed to delete documents from PostgreSQL"
+
+
+@app.command()
+def similarity_search(
+    query: str = typer.Option(
+        "禅モード",
+        "--query",
+        "-q",
+        help="Query to search in the Cosmos DB index",
+    ),
+    k: int = typer.Option(
+        5,
+        "--k",
+        "-k",
+        help="Number of results to return from the similarity search",
+    ),
+    verbose: bool = typer.Option(
+        False,
+        "--verbose",
+        "-v",
+        help="Enable verbose output",
+    ),
+):
+    # Set up logging
+    if verbose:
+        logger.setLevel(logging.DEBUG)
+
+    logger.info(f"Searching Cosmos DB with query: {query}")
+
+    # Perform similarity search
+    postgresql_client = PostgreSQLClientWrapper()
+    documents = postgresql_client.similarity_search(
+        query=query,
+        k=k,  # Number of results to return
+    )
+    logger.info(f"Found {len(documents)} results for query: {query}")
+
+    # Log the results
+    for i, document in enumerate(documents, start=1):
+        logger.debug("-" * 40)
+        logger.debug(f"#{i}: {document.model_dump_json(indent=2)}")
+
+
+if __name__ == "__main__":
+    load_dotenv(
+        override=True,
+        verbose=True,
+    )
+    app()

template_langgraph/tools/postgresql_tool.py

@@ -0,0 +1,60 @@
+from functools import lru_cache
+
+from langchain_core.documents import Document
+from langchain_postgres import PGEngine, PGVectorStore
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
+
+
+class Settings(BaseSettings):
+    postgresql_user: str = "user"
+    postgresql_password: str = "password"
+    postgresql_host: str = "localhost"
+    postgresql_port: str = "5432"
+    postgresql_database: str = "db"
+    postgresql_table_name: str = "reports_kabuto_vectors"
+    postgresql_id_column: str = "id"
+    postgresql_content_column: str = "user_message"
+    postgresql_embedding_column: str = "embedding"
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_ignore_empty=True,
+        extra="ignore",
+    )
+
+
+@lru_cache
+def get_postgresql_settings() -> Settings:
+    """Get PostgreSQL Database settings."""
+    return Settings()
+
+
+class PostgreSQLClientWrapper:
+    def __init__(
+        self,
+        settings: Settings = None,
+    ):
+        if settings is None:
+            settings = get_postgresql_settings()
+
+        self.vector_store = PGVectorStore.create_sync(
+            engine=PGEngine.from_connection_string(
+                url=f"postgresql+psycopg://{settings.postgresql_user}:{settings.postgresql_password}@{settings.postgresql_host}:{settings.postgresql_port}/{settings.postgresql_database}"
+            ),
+            table_name=settings.postgresql_table_name,
+            embedding_service=AzureOpenAiWrapper().embedding_model,
+            id_column=settings.postgresql_id_column,
+            content_column=settings.postgresql_content_column,
+            embedding_column=settings.postgresql_embedding_column,
+        )
+
+    def add_documents(
+        self,
+        documents: list[Document],
+    ) -> list[str]:
+        """Add documents to a Cosmos DB container."""
+        return self.vector_store.add_documents(
+            documents=documents,
+        )