add AI Search tool

Author: ks6088ts

docs/references.md

@@ -21,6 +21,7 @@
 - [CSVLoader](https://python.langchain.com/docs/how_to/document_loader_csv/)
 - [Qdrant](https://github.com/qdrant/qdrant)
 - [Azure Cosmos DB No SQL](https://python.langchain.com/docs/integrations/vectorstores/azure_cosmos_db_no_sql/)
+- [Azure AI Search](https://python.langchain.com/docs/integrations/vectorstores/azuresearch/)
 
 ### Services
 

template_langgraph/tools/ai_search_tool.py

@@ -2,6 +2,8 @@
 
 from langchain_community.vectorstores.azuresearch import AzureSearch
 from langchain_core.documents import Document
+from langchain_core.tools import tool
+from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
@@ -58,3 +60,46 @@ def similarity_search(
             query=query,
             k=k,  # Number of results to return
         )
+
+
+class AiSearchInput(BaseModel):
+    query: str = Field(
+        default="禅モード",
+        description="Query to search in the AI Search index",
+    )
+    k: int = Field(
+        default=5,
+        description="Number of results to return from the similarity search",
+    )
+
+
+class AiSearchOutput(BaseModel):
+    content: str = Field(description="Content of the document")
+    id: str = Field(description="ID of the document")
+
+
+@tool(args_schema=AiSearchInput)
+def search_ai_search(query: str, k: int = 5) -> list[AiSearchOutput]:
+    """Search for similar documents in AI Search index.
+
+    Args:
+        query: The search query string
+        k: Number of results to return (default: 5)
+
+    Returns:
+        AiSearchOutput: A Pydantic model containing the search results
+    """
+    wrapper = AiSearchClientWrapper()
+    documents = wrapper.similarity_search(
+        query=query,
+        k=k,
+    )
+    outputs = []
+    for document in documents:
+        outputs.append(
+            {
+                "content": document.page_content,
+                "id": document.id,
+            }
+        )
+    return outputs

template_langgraph/tools/common.py

@@ -1,5 +1,6 @@
 from template_langgraph.llms.azure_openais import AzureOpenAiWrapper
 from template_langgraph.loggers import get_logger
+from template_langgraph.tools.ai_search_tool import search_ai_search
 from template_langgraph.tools.cosmosdb_tool import search_cosmosdb
 from template_langgraph.tools.dify_tool import run_dify_workflow
 from template_langgraph.tools.elasticsearch_tool import search_elasticsearch
@@ -18,6 +19,7 @@ def get_default_tools():
         logger.error(f"Error occurred while getting SQL database tools: {e}")
         sql_database_tools = []
     return [
+        search_ai_search,
         search_cosmosdb,
         run_dify_workflow,
         search_qdrant,