Ryan Chen
32020a6c60
- Update llm.py to use OpenAI client with custom base_url for llama-server - Update agents.py to use ChatOpenAI instead of ChatOllama - Remove unused ollama imports from main.py, chunker.py, query.py - Add LLAMA_SERVER_URL and LLAMA_MODEL_NAME env vars - Remove ollama and langchain-ollama dependencies Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
import os
|
|
from typing import cast
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain.chat_models import BaseChatModel
|
|
from langchain.tools import tool
|
|
from langchain_openai import ChatOpenAI
|
|
from tavily import AsyncTavilyClient
|
|
|
|
from blueprints.rag.logic import query_vector_store
|
|
|
|
# Configure LLM with llama-server or OpenAI fallback
|
|
llama_url = os.getenv("LLAMA_SERVER_URL")
|
|
if llama_url:
|
|
llama_chat = ChatOpenAI(
|
|
base_url=llama_url,
|
|
api_key="not-needed",
|
|
model=os.getenv("LLAMA_MODEL_NAME", "llama-3.1-8b-instruct"),
|
|
)
|
|
else:
|
|
llama_chat = None
|
|
|
|
openai_fallback = ChatOpenAI(model="gpt-5-mini")
|
|
model_with_fallback = cast(
|
|
BaseChatModel,
|
|
llama_chat.with_fallbacks([openai_fallback]) if llama_chat else openai_fallback,
|
|
)
|
|
client = AsyncTavilyClient(os.getenv("TAVILY_KEY"), "")
|
|
|
|
|
|
@tool
|
|
async def web_search(query: str) -> str:
|
|
"""Search the web for current information using Tavily.
|
|
|
|
Use this tool when you need to:
|
|
- Find current information not in the knowledge base
|
|
- Look up recent events, news, or updates
|
|
- Verify facts or get additional context
|
|
- Search for information outside of Simba's documents
|
|
|
|
Args:
|
|
query: The search query to look up on the web
|
|
|
|
Returns:
|
|
Search results from the web with titles, content, and source URLs
|
|
"""
|
|
response = await client.search(query=query, search_depth="basic")
|
|
results = response.get("results", [])
|
|
|
|
if not results:
|
|
return "No results found for the query."
|
|
|
|
formatted = "\n\n".join(
|
|
[
|
|
f"**{result['title']}**\n{result['content']}\nSource: {result['url']}"
|
|
for result in results[:5]
|
|
]
|
|
)
|
|
return formatted
|
|
|
|
|
|
@tool(response_format="content_and_artifact")
|
|
async def simba_search(query: str):
|
|
"""Search through Simba's medical records, veterinary documents, and personal information.
|
|
|
|
Use this tool whenever the user asks questions about:
|
|
- Simba's health history, medical records, or veterinary visits
|
|
- Medications, treatments, or diagnoses
|
|
- Weight, diet, or physical characteristics over time
|
|
- Veterinary recommendations or advice
|
|
- Ryan's (the owner's) information related to Simba
|
|
- Any factual information that would be found in documents
|
|
|
|
Args:
|
|
query: The user's question or information need about Simba
|
|
|
|
Returns:
|
|
Relevant information from Simba's documents
|
|
"""
|
|
print(f"[SIMBA SEARCH] Tool called with query: {query}")
|
|
serialized, docs = await query_vector_store(query=query)
|
|
print(f"[SIMBA SEARCH] Found {len(docs)} documents")
|
|
print(f"[SIMBA SEARCH] Serialized result length: {len(serialized)}")
|
|
print(f"[SIMBA SEARCH] First 200 chars: {serialized[:200]}")
|
|
return serialized, docs
|
|
|
|
|
|
main_agent = create_agent(model=model_with_fallback, tools=[simba_search, web_search])
|