File size: 2,274 Bytes
c6d67ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | import os
from langchain_core.tools import tool
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.messages import SystemMessage, HumanMessage
import functools
@functools.lru_cache(maxsize=1)
def get_cached_embeddings():
return HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
def get_10k_vector_db(ticker: str) -> Chroma:
"""Loads a pre-computed 10-K Chroma Vector Database from disk."""
ticker = ticker.upper()
persist_directory = f"./chroma_db/{ticker}_10k"
if not os.path.exists(persist_directory):
raise FileNotFoundError(
f"10-K data for {ticker} has not been ingested. "
f"Please run the ingestion pipeline: `python ingest.py --tickers {ticker}`"
)
# Using cached embeddings to prevent massive memory slowdowns on agent loops
embeddings = get_cached_embeddings()
return Chroma(
persist_directory=persist_directory,
embedding_function=embeddings
)
@tool
def search_10k_filings(ticker: str, query: str, llm=None) -> str:
"""Searches 10-K and returns a CONCISE summary of findings."""
try:
db = get_10k_vector_db(ticker)
results = db.similarity_search(query, k=2)
if not results:
return f"No info found for {query}."
# Combine the text
context = "\n".join([doc.page_content for doc in results])
# We can use the LLM to 'pre-process' the data so the Supervisor stays clean
# Note: You'll need to pass the 'llm' object into this tool or initialize a local one
if llm:
response = llm.invoke([
SystemMessage(content="You are a helpful assistant."),
HumanMessage(content=f"Summarize the following 10-K findings for {ticker} regarding {query}:\n\n{context}")
])
return response.content
else:
return f"SUMMARY OF 10-K FINDINGS FOR {ticker} ({query}):\n\n{context}"
except Exception as e:
return f"Error: {str(e)}"
if __name__ == "__main__":
# Test using .invoke() as we discussed for sec_tools
print(search_10k_filings.invoke({"ticker": "TSLA", "query": "marketing risks"})) |