Implemene Azure search tool

requirements.txt

@@ -23,3 +23,6 @@ yfinance==0.2.36
 transformers==4.40.2
 langchain==0.1.9
 accelerate==0.27.2
+azure-search-documents==11.6.0b1
+azure-identity==1.16.0
+

src/config.py

@@ -5,7 +5,13 @@ class Config(BaseSettings):
     hf_token: str = Field(...)
     hf_model: str = Field("InvestmentResearchAI/LLM-ADE-dev")
     headless: bool = Field(False, description="Run in headless mode.")
-    
+
+    az_search_endpoint: str = Field("https://analysis-bank.search.windows.net")
+    az_search_api_key: str = Field(...)
+    az_search_idx_name: str = Field("analysis-index")
+    az_search_top_k: int = Field(2, description="Max number of results to retrun")
+    az_search_min_score: float = Field(12.0, description="Only results above this confidence score is used")
+
     chat_template: str = Field("chatml", description="Chat template for prompt formatting")
     num_fewshot: int | None = Field(None, description="Option to use json mode examples")
     load_in_4bit: str = Field("False", description="Option to load in 4bit with bitsandbytes")

src/functions.py

@@ -11,6 +11,38 @@ from utils import inference_logger
 from langchain.tools import tool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 
+from azure.core.credentials import AzureKeyCredential
+from azure.search.documents import SearchClient
+
+
+az_creds = AzureKeyCredential(config.az_search_api_key)
+az_search_client = SearchClient(config.az_search_endpoint, config.az_search_idx_name, az_creds)
+
+@tool
+def get_company_analysis(query: str) -> dict:
+    """
+    Searches through your database of company and crypto analysis, retrieves top 2 
+    pieces of analysis relevant to your query.
+    
+    Args:
+        query (str): The search query
+    Returns:
+        list: A list of dictionaries containing the pieces of analysis.
+    """
+    results = az_search_client.search(
+            query_type="simple", 
+            search_text=query, 
+            select="title,content", 
+            include_total_count=True, 
+            top=config.az_search_top_k
+        )
+    
+    output = []
+    for x in results:
+        if x["@search.score"] >= config.az_search_min_score:
+            output.append({"title": x["title"], "content": x["content"]})
+    return output
+
 @tool
 def google_search_and_scrape(query: str) -> dict:
     """
@@ -246,6 +278,7 @@ def get_company_profile(symbol: str) -> dict:
 
 def get_openai_tools() -> List[dict]:
     functions = [
+        get_company_analysis,  
         google_search_and_scrape,
         get_current_stock_price,
         get_company_news,