Update get_analysis to better version

ex.env

@@ -0,0 +1,4 @@
+HF_TOKEN=
+OLLAMA_MODEL=
+AZ_SEARCH_API_KEY=
+AZURE_OPENAI_API_KEY=

requirements.txt

@@ -26,3 +26,4 @@ accelerate==0.27.2
 azure-search-documents==11.6.0b1
 azure-identity==1.16.0
 loguru==0.7.2
+openai==1.30.1

src/config.py

@@ -17,9 +17,11 @@ class Config(BaseSettings):
 
     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_idx_name: str = Field("analysis-index-2024-05-19")
     az_search_top_k: int = Field(4, description="Max number of results to retrun")
-    az_search_min_score: float = Field(9.0, description="Only results above this confidence score is used")
+    
+    azure_openai_api_key: str = Field(...)
+    azure_openai_endpoint: str = Field("https://irai-openai-eastus.openai.azure.com/")
 
     chat_template: str = Field("chatml", description="Chat template for prompt formatting")
     num_fewshot: int | None = Field(None, description="Option to use json mode examples")

src/functions.py

@@ -1,24 +1,31 @@
 import re
-import inspect
 import requests
 import pandas as pd
 import yfinance as yf
 import concurrent.futures
+from time import time
 from datetime import datetime
 
 from typing import List
 from bs4 import BeautifulSoup
 from logger import logger
+from openai import AzureOpenAI
 from langchain.tools import tool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from config import config
 
 from azure.core.credentials import AzureKeyCredential
 from azure.search.documents import SearchClient
+from azure.search.documents.models import VectorizedQuery
 
 
 az_creds = AzureKeyCredential(config.az_search_api_key)
 az_search_client = SearchClient(config.az_search_endpoint, config.az_search_idx_name, az_creds)
+openai_client = AzureOpenAI(
+    azure_endpoint=config.azure_openai_endpoint, 
+    api_key=config.azure_openai_api_key,
+    api_version="2024-02-01"
+)
 
 @tool
 def get_analysis(query: str) -> dict:
@@ -32,25 +39,29 @@ def get_analysis(query: str) -> dict:
     Returns:
         list: A list of dictionaries containing the pieces of analysis.
     """
+    start_time = time() 
+    
+    embed_model = "default-large-embeddings"
+    
+    vec = openai_client.embeddings.create(input=[query], model=embed_model).data[0].embedding
+    vector_query = VectorizedQuery(vector=vec, k_nearest_neighbors=config.az_search_top_k * 2, fields="vector")
     results = az_search_client.search(
-            query_type="semantic", 
-            search_text=query, 
-            select="title,content,asset_name,write_date", 
-            include_total_count=True, 
-            top=config.az_search_top_k,
-            semantic_configuration_name="basic-keywords",
-            vector_queries=None, # Docs are too semantically similar, disable for now
-        )
+        search_text="*",
+        vector_queries=[vector_query],
+        select=["date", "popularity", "sequence", "context", "securities"],
+        order_by=["securities desc", "date desc", "popularity desc"],
+        top=config.az_search_top_k,
+    )
     
     output = []
     for x in results:
-        if x["@search.score"] >= config.az_search_min_score:
-            output.append({
-                "security": x["asset_name"],
-                "date written": datetime.strptime(x["write_date"], "%Y%m%d").date(),
-                "title": x["title"], 
-                "content": x["content"]
-            })
+        output.append({
+            "securities": x["securities"],
+            "date written": x["date"].split("T")[0],
+            "summary": x["context"], 
+            "content": x["sequence"],
+        })
+    print(f"Search took {time() - start_time:.2f} seconds\n---")
     return output
 
 @tool

tests/test_functions.py

@@ -0,0 +1,13 @@
+# Tests the functions themselves, not the function calling
+
+from pprint import pprint
+from functions import get_analysis
+
+def test_get_analysis():
+    query = "How is MSTR doing?"
+    output = get_analysis(query)
+    
+    pprint(output)
+    assert len(output) != 0
+    assert "MSTR" in output[0]['securities']
+