add a CLI to run AI Search operation

Author: ks6088ts

.env.template

@@ -45,6 +45,11 @@ COSMOSDB_PARTITION_KEY="/id"
 SQL_DATABASE_URI="sqlite:///template_langgraph.db"
 # SQL_DATABASE_URI="postgresql://user:password@localhost:5432/db"
 
+# Azure AI Search Settings
+AI_SEARCH_ENDPOINT="https://xxx.search.windows.net/"
+AI_SEARCH_KEY="xxx"
+AI_SEARCH_INDEX_NAME="kabuto"
+
 # ---------
 # Utilities
 # ---------

scripts/ai_search_operator.py

@@ -0,0 +1,102 @@
+import logging
+
+import typer
+from dotenv import load_dotenv
+
+from template_langgraph.loggers import get_logger
+from template_langgraph.tools.ai_search_tool import AiSearchClientWrapper
+from template_langgraph.utilities.csv_loaders import CsvLoaderWrapper
+from template_langgraph.utilities.pdf_loaders import PdfLoaderWrapper
+
+# Initialize the Typer application
+app = typer.Typer(
+    add_completion=False,
+    help="AI Search 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 AI Search
+    ai_search_client = AiSearchClientWrapper()
+    ids = ai_search_client.add_documents(
+        documents=documents,
+    )
+    logger.info(f"Added {len(ids)} documents to AI Search.")
+    for id in ids:
+        logger.debug(f"Added document ID: {id}")
+
+
+@app.command()
+def similarity_search(
+    query: str = typer.Option(
+        "禅モード",
+        "--query",
+        "-q",
+        help="Query to search in the AI Search 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 AI Search with query: {query}")
+
+    # Perform similarity search
+    ai_search_client = AiSearchClientWrapper()
+    documents = ai_search_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/ai_search_tool.py

@@ -0,0 +1,60 @@
+from functools import lru_cache
+
+from langchain_community.vectorstores.azuresearch import AzureSearch
+from langchain_core.documents import Document
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
+
+
+class Settings(BaseSettings):
+    ai_search_key: str = "<your-ai-search-key>"
+    ai_search_endpoint: str = "<your-ai-search-endpoint>"
+    ai_search_index_name: str = "<your-ai-index-name>"
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_ignore_empty=True,
+        extra="ignore",
+    )
+
+
+@lru_cache
+def get_ai_search_settings() -> Settings:
+    """Get AI Search settings."""
+    return Settings()
+
+
+class AiSearchClientWrapper:
+    def __init__(
+        self,
+        settings: Settings = None,
+    ):
+        if settings is None:
+            settings = get_ai_search_settings()
+        self.vector_store: AzureSearch = AzureSearch(
+            azure_search_endpoint=settings.ai_search_endpoint,
+            azure_search_key=settings.ai_search_key,
+            index_name=settings.ai_search_index_name,
+            embedding_function=AzureOpenAiWrapper().embedding_model.embed_query,
+        )
+
+    def add_documents(
+        self,
+        documents: list[Document],
+    ) -> list[str]:
+        """Add documents to a Cosmos DB container."""
+        return self.vector_store.add_documents(
+            documents=documents,
+        )
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 5,
+    ) -> list[Document]:
+        """Perform a similarity search in the Cosmos DB index."""
+        return self.vector_store.similarity_search(
+            query=query,
+            k=k,  # Number of results to return
+        )