| """The four specialist tools, ported to the local model.
|
|
|
| Behaviour matches `agent_tools.py`; only the model behind each call changed.
|
| Print statements are gone — node progress now travels over the event stream.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import asyncio
|
| import sqlite3
|
| from typing import Any, Dict, List
|
|
|
| import pandas as pd
|
| from langchain_core.tools import tool
|
| from pydantic import BaseModel, Field
|
|
|
| from ..config import get_settings
|
| from .embeddings import assert_dimension_matches, get_embeddings
|
| from .llm import chat_model, structured_model
|
| from .resources import collection_dimension, get_qdrant, get_reranker, get_sql_database
|
|
|
| _QUERY_REWRITE = (
|
| 'Rewrite the following query for a semantic search system. Respond with ONLY '
|
| 'the rewritten query, no preamble.\nQuery: "{query}"\nRewritten:'
|
| )
|
|
|
|
|
| async def optimize_query(query: str) -> str:
|
| """Rewrite a user query for retrieval.
|
|
|
| Falls back to the original query on failure: a bad rewrite should degrade
|
| retrieval, not fail the request.
|
| """
|
| try:
|
| response = await chat_model(temperature=0.0, max_tokens=128).ainvoke(
|
| _QUERY_REWRITE.format(query=query)
|
| )
|
| rewritten = (response.content or "").strip().strip('"')
|
| return rewritten or query
|
| except Exception:
|
| return query
|
|
|
|
|
| @tool
|
| async def librarian_rag_tool(query: str) -> List[Dict[str, Any]]:
|
| """Retrieves deep, contextual information from Alphabet's financial filings."""
|
| settings = get_settings()
|
| assert_dimension_matches(collection_dimension())
|
|
|
| optimized = await optimize_query(query)
|
| embedding = await asyncio.to_thread(get_embeddings().embed_query, optimized)
|
|
|
| hits = await asyncio.to_thread(
|
| lambda: get_qdrant().query_points(
|
| collection_name=settings.collection_name,
|
| query=embedding,
|
| limit=settings.retrieval_candidates,
|
| with_payload=True,
|
| ).points
|
| )
|
| if not hits:
|
| return []
|
|
|
|
|
| pairs = [[optimized, hit.payload.get("content", "")] for hit in hits]
|
| scores = await asyncio.to_thread(get_reranker().predict, pairs)
|
|
|
| ranked = sorted(zip(hits, scores), key=lambda pair: pair[1], reverse=True)
|
| return [
|
| {
|
| "source": hit.payload.get("source"),
|
| "content": hit.payload.get("content"),
|
| "summary": hit.payload.get("summary"),
|
| "rerank_score": float(score),
|
| }
|
| for hit, score in ranked[: settings.retrieval_top_k]
|
| ]
|
|
|
|
|
| class _SqlQuery(BaseModel):
|
| sql: str = Field(description="A single read-only SQLite SELECT statement.")
|
|
|
|
|
| _SQL_PROMPT = """Write one SQLite SELECT query answering the question.
|
|
|
| Schema:
|
| {schema}
|
|
|
| Rules:
|
| - SELECT only. No INSERT, UPDATE, DELETE, DROP or PRAGMA.
|
| - Use only the columns shown above.
|
| - SQLite syntax: no YEAR(), no TOP. Use LIMIT.
|
| - Return the query only.
|
|
|
| Question: {question}
|
| """
|
|
|
|
|
|
|
|
|
| _FORBIDDEN_SQL = (
|
| "insert", "update", "delete", "drop", "alter", "create",
|
| "attach", "detach", "pragma", "replace", "vacuum",
|
| )
|
|
|
|
|
| def _run_select(sql: str) -> str:
|
| settings = get_settings()
|
| stripped = sql.strip().rstrip(";").strip()
|
| lowered = stripped.lower()
|
|
|
| if not (lowered.startswith("select") or lowered.startswith("with")):
|
| return f"Refused: only SELECT queries are allowed, got: {stripped[:120]}"
|
| if any(word in lowered.split() for word in _FORBIDDEN_SQL):
|
| return f"Refused: statement contains a write operation: {stripped[:120]}"
|
| if ";" in stripped:
|
| return "Refused: multiple statements are not allowed."
|
|
|
| with sqlite3.connect(f"file:{settings.sqlite_path}?mode=ro", uri=True) as conn:
|
| df = pd.read_sql_query(stripped, conn)
|
|
|
| if df.empty:
|
| return f"Query returned no rows.\nSQL: {stripped}"
|
| return f"SQL: {stripped}\n\n{df.to_string(index=False, max_rows=40)}"
|
|
|
|
|
| @tool
|
| async def analyst_sql_tool(query: str) -> str:
|
| """Answers specific questions about revenue, net income, etc., from a SQL database."""
|
|
|
|
|
|
|
|
|
| try:
|
| schema = await asyncio.to_thread(get_sql_database().get_table_info)
|
| except Exception as exc:
|
| return f"Could not read database schema: {type(exc).__name__}: {exc}"
|
|
|
| try:
|
| generated = await structured_model(_SqlQuery).ainvoke(
|
| _SQL_PROMPT.format(schema=schema, question=query)
|
| )
|
| except Exception as exc:
|
| return f"Could not generate SQL: {type(exc).__name__}: {exc}"
|
|
|
| try:
|
| return await asyncio.to_thread(_run_select, generated.sql)
|
| except Exception as exc:
|
| return f"SQL failed: {type(exc).__name__}: {exc}\nSQL: {generated.sql}"
|
|
|
|
|
| def _trend_sync(_: str) -> str:
|
| settings = get_settings()
|
| with sqlite3.connect(settings.sqlite_path) as conn:
|
| df = pd.read_sql_query(
|
| f"SELECT * FROM {settings.table_name} ORDER BY year, quarter", conn
|
| )
|
| if df.empty:
|
| return "No rows in the financial summary table."
|
|
|
| df["period"] = df["year"].astype(str) + "-" + df["quarter"].astype(str)
|
| df = df.set_index("period")
|
| metric = "revenue_usd_billions"
|
|
|
| qoq = df[metric].pct_change()
|
| yoy = df[metric].pct_change(4)
|
| return (
|
| f"Analysis of {metric} from {df.index[0]} to {df.index[-1]}: "
|
| f"${df[metric].iloc[0]:.1f}B -> ${df[metric].iloc[-1]:.1f}B. "
|
| f"Latest QoQ {qoq.iloc[-1]:.1%}, latest YoY {yoy.iloc[-1]:.1%}."
|
| )
|
|
|
|
|
| @tool
|
| async def analyst_trend_tool(query: str) -> str:
|
| """Analyzes financial trends like QoQ or YoY growth from the SQL database."""
|
| return await asyncio.to_thread(_trend_sync, query)
|
|
|
|
|
| @tool
|
| async def scout_web_search_tool(query: str) -> str:
|
| """Searches the web for real-time information like stock prices or recent news."""
|
| from langchain_community.tools import DuckDuckGoSearchRun
|
|
|
| try:
|
| return await DuckDuckGoSearchRun().arun(query)
|
| except Exception as exc:
|
| return f"Web search failed: {type(exc).__name__}: {exc}"
|
|
|
|
|
| ALL_TOOLS = [
|
| librarian_rag_tool,
|
| analyst_sql_tool,
|
| analyst_trend_tool,
|
| scout_web_search_tool,
|
| ]
|
| TOOL_MAP = {t.name: t for t in ALL_TOOLS}
|
|
|
