File size: 4,740 Bytes
b2931f4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
"""Agent tools β€” provider-neutral registry.

Each tool is a plain Python function (independently testable) plus a `ToolSpec`
describing its name, when-to-use text, and JSON-schema parameters. The spec is
deliberately NOT in any vendor's tool format: Decision 16 adapts these into
LangChain/LangGraph tools (which bind to Gemini or Claude alike), keeping the
provider seam from Decision 14 intact. Hardwiring Anthropic's tool_use shape
here β€” as the original handoff assumed β€” would have undone that.

`dispatch(name, args)` is the single call site the agent loop uses to run a
tool by name; it returns the tool's dict result unchanged.

Three tools, by design β€” see [[project_finrag_overview]]:
  calculator       β€” arithmetic, safe-eval        (pure fn)
  lookup_citation  β€” re-fetch a chunk from Qdrant  (read-only)
  sql_query        β€” NL β†’ SQL over DuckDB facts     (sub-LLM; added next)
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable

from finrag.tools.calculator import calculator
from finrag.tools.citation import lookup_citation
from finrag.tools.sql import sql_query


@dataclass(frozen=True)
class ToolSpec:
    name: str
    description: str
    # JSON-schema "object" describing the args. Both Gemini and Anthropic
    # accept this shape (modulo a thin adapter), as does LangChain.
    parameters: dict[str, Any]
    fn: Callable[..., dict[str, Any]]


TOOL_SPECS: list[ToolSpec] = [
    ToolSpec(
        name="calculator",
        description=(
            "Evaluate an arithmetic expression over numeric literals "
            "(+ - * / // % ** and parentheses). Use for growth rates, margins, "
            "sums, and ratios instead of doing mental math. Extract the numbers "
            "from context first, then pass an expression like "
            "'(383285 - 394328) / 394328 * 100'."
        ),
        parameters={
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Arithmetic expression over numeric literals only.",
                }
            },
            "required": ["expression"],
        },
        fn=calculator,
    ),
    ToolSpec(
        name="lookup_citation",
        description=(
            "Re-fetch the full text and provenance of a single retrieved chunk "
            "by its chunk_id. Use when you need to quote an exact figure or "
            "re-read a chunk you cited earlier."
        ),
        parameters={
            "type": "object",
            "properties": {
                "chunk_id": {
                    "type": "string",
                    "description": (
                        "The chunk_id of a previously retrieved chunk β€” the exact "
                        "value shown as (id=...) in that chunk's header. Do not "
                        "use the [N] anchor or invent an id."
                    ),
                }
            },
            "required": ["chunk_id"],
        },
        fn=lookup_citation,
    ),
    ToolSpec(
        name="sql_query",
        description=(
            "Query exact financial figures (revenue, net income, R&D, margins, "
            "multi-year or cross-company comparisons) from the structured "
            "financial_facts database. Pass a natural-language description of "
            "the numbers you need; SQL is generated and run for you. Prefer this "
            "over reading figures out of text chunks when precision matters."
        ),
        parameters={
            "type": "object",
            "properties": {
                "natural_language": {
                    "type": "string",
                    "description": "Plain-language description of the figures to fetch.",
                }
            },
            "required": ["natural_language"],
        },
        fn=sql_query,
    ),
]

# name β†’ spec, for O(1) dispatch.
TOOL_REGISTRY: dict[str, ToolSpec] = {spec.name: spec for spec in TOOL_SPECS}


def dispatch(name: str, args: dict[str, Any]) -> dict[str, Any]:
    """Run tool `name` with keyword `args`. Returns the tool's dict result.

    Unknown tool names return an error dict (not a raise) so a hallucinated
    tool call degrades gracefully inside the agent loop.
    """
    spec = TOOL_REGISTRY.get(name)
    if spec is None:
        return {"error": f"Unknown tool {name!r}. Available: {list(TOOL_REGISTRY)}"}
    try:
        return spec.fn(**args)
    except TypeError as e:
        # Wrong/missing args from the model β€” surface as data, not a crash.
        return {"error": f"Bad arguments for {name!r}: {e}"}


__all__ = ["ToolSpec", "TOOL_SPECS", "TOOL_REGISTRY", "dispatch"]