Merge pull request #256 from krishsharma-code/feat/rag-calculator-tool

backend/app/rag/agent.py

@@ -11,6 +11,7 @@ from app.config import get_settings
 from app.rag.retriever import retrieve
 from app.rag.graph_retriever import get_entity_context
 from app.rag.prompts import SYSTEM_PROMPT, RAG_PROMPT_TEMPLATE, GREETING_PROMPT
+from app.rag.tools import TOOL_PROMPT, TOOLS, execute_tool
 from app.rag.tracing import trace_function
 
 logger = logging.getLogger(__name__)
@@ -23,6 +24,34 @@ def get_llm_client(hf_token: Optional[str] = None) -> InferenceClient:
     )
 
 
+def _execute_tools_if_requested(client: InferenceClient, messages: list[dict[str, Any]]) -> Any:
+    """Run the LLM and execute any tool call responses until the final answer is produced."""
+    for _ in range(3):
+        response = client.chat_completion(
+            messages=messages,
+            model=settings.LLM_MODEL,
+            max_tokens=settings.LLM_MAX_NEW_TOKENS,
+            temperature=settings.LLM_TEMPERATURE,
+            tools=TOOLS,
+            tool_prompt=TOOL_PROMPT,
+        )
+
+        choice = response.choices[0]
+        tool_calls = getattr(choice.message, "tool_calls", None)
+        if not tool_calls:
+            return response
+
+        tool_call = tool_calls[0]
+        tool_name = tool_call.function.name
+        tool_args = json.loads(tool_call.function.arguments)
+        tool_result = execute_tool(tool_name, tool_args)
+
+        messages.append({"role": "tool", "name": tool_name, "content": tool_result})
+
+    # If tools are still requested after several rounds, return the latest response anyway.
+    return response
+
+
 def is_greeting(question: str) -> bool:
     """Detect if the question is a casual greeting rather than a document query."""
     greetings = {
@@ -141,12 +170,7 @@ def generate_answer(
     # ── Generate answer ──────────────────────────────
     # STAGE 3: Send prompt to HuggingFace Inference API and get the generated answer
     try:
-        response = client.chat_completion(
-            messages=messages,
-            model=settings.LLM_MODEL,
-            max_tokens=settings.LLM_MAX_NEW_TOKENS,
-            temperature=settings.LLM_TEMPERATURE,
-        )
+        response = _execute_tools_if_requested(client, messages)
         if response.choices:
             answer = response.choices[0].message.content.strip()
         else:
@@ -257,15 +281,17 @@ def generate_answer_stream(
     user_content = RAG_PROMPT_TEMPLATE.format(context=context, question=question)
     messages = _chat_messages(SYSTEM_PROMPT, user_content)
 
-    # ── Stream answer tokens ─────────────────────────
-    # STAGE 3: Stream tokens from HuggingFace Inference API → forward each as an SSE 'token' event
+    # Resolve tool calls before streaming, then stream the final answer.
     try:
+        _execute_tools_if_requested(client, messages)
         stream = client.chat_completion(
             messages=messages,
             model=settings.LLM_MODEL,
             max_tokens=settings.LLM_MAX_NEW_TOKENS,
             temperature=settings.LLM_TEMPERATURE,
             stream=True,
+            tools=TOOLS,
+            tool_prompt=TOOL_PROMPT,
         )
         for chunk in stream:
             if chunk.choices:

backend/app/rag/prompts.py

@@ -12,6 +12,7 @@ IMPORTANT RULES:
 4. Be precise, clear, and well-structured in your responses.
 5. Use bullet points and formatting when listing multiple items.
 6. For numerical data or key facts, quote the relevant text directly.
+7. If a question requires arithmetic calculations, use the registered calculator tool instead of guessing or estimating.
 
 FORMATTING:
 - Use **bold** for key terms and important findings

backend/app/rag/tools.py

@@ -0,0 +1,116 @@
+"""Agent tools for the PDF Assistant RAG backend."""
+
+import ast
+import operator as op
+from typing import Any
+
+from huggingface_hub.inference._generated.types.chat_completion import (
+    ChatCompletionInputFunctionDefinition,
+    ChatCompletionInputTool,
+)
+
+_ALLOWED_OPERATORS = {
+    ast.Add: op.add,
+    ast.Sub: op.sub,
+    ast.Mult: op.mul,
+    ast.Div: op.truediv,
+    ast.FloorDiv: op.floordiv,
+    ast.Mod: op.mod,
+    ast.Pow: op.pow,
+    ast.USub: op.neg,
+    ast.UAdd: op.pos,
+}
+
+
+def _evaluate_ast(node: ast.AST) -> float:
+    if isinstance(node, ast.Expression):
+        return _evaluate_ast(node.body)
+
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, (int, float)):
+            return float(node.value)
+        raise ValueError("Only numeric values are allowed in calculator expressions.")
+
+    if isinstance(node, ast.BinOp):
+        left = _evaluate_ast(node.left)
+        right = _evaluate_ast(node.right)
+        operator = type(node.op)
+        if operator not in _ALLOWED_OPERATORS:
+            raise ValueError(f"Operator {operator.__name__} is not allowed.")
+        return _ALLOWED_OPERATORS[operator](left, right)
+
+    if isinstance(node, ast.UnaryOp):
+        operator = type(node.op)
+        if operator not in _ALLOWED_OPERATORS:
+            raise ValueError(f"Operator {operator.__name__} is not allowed.")
+        operand = _evaluate_ast(node.operand)
+        return _ALLOWED_OPERATORS[operator](operand)
+
+    raise ValueError("Unsupported expression in calculator tool.")
+
+
+def calculate_expression(expression: str) -> str:
+    """Safely evaluate a simple arithmetic expression.
+
+    This tool only permits numeric literals and arithmetic operators.
+    It does not execute arbitrary code.
+    """
+    try:
+        parsed = ast.parse(expression, mode="eval")
+    except SyntaxError as exc:
+        raise ValueError(f"Invalid calculator expression: {exc}") from exc
+
+    if not isinstance(parsed, ast.Expression):
+        raise ValueError("Expression must be a single arithmetic expression.")
+
+    result = _evaluate_ast(parsed)
+
+    if result.is_integer():
+        return str(int(result))
+
+    return str(result)
+
+
+def execute_tool(name: str, arguments: dict[str, Any]) -> str:
+    """Execute a registered tool by name."""
+    if name != "calculator":
+        raise ValueError(f"Unknown tool: {name}")
+
+    expression = arguments.get("expression")
+    if not isinstance(expression, str) or not expression.strip():
+        raise ValueError("The calculator tool requires a non-empty 'expression' string.")
+
+    return calculate_expression(expression)
+
+
+CALCULATOR_TOOL = ChatCompletionInputTool(
+    function=ChatCompletionInputFunctionDefinition(
+        name="calculator",
+        description=(
+            "Safely evaluate a numeric arithmetic expression for financial calculations. "
+            "Use only numeric values and arithmetic operators like +, -, *, /, %, //, and **."
+        ),
+        parameters={
+            "type": "object",
+            "properties": {
+                "expression": {
+                    "type": "string",
+                    "description": (
+                        "A valid arithmetic expression to evaluate, for example '1000 - 250' or "
+                        "'(revenue - expenses) * 0.2'."
+                    ),
+                }
+            },
+            "required": ["expression"],
+        },
+    ),
+    type="tool",
+)
+
+TOOL_PROMPT = (
+    "Use the calculator tool for all numeric arithmetic operations in the user query. "
+    "The tool accepts a single 'expression' field and returns the evaluated numeric result. "
+    "Do not attempt to compute arithmetic without the tool."
+)
+
+TOOLS = [CALCULATOR_TOOL]

backend/tests/test_rag_tools.py

@@ -0,0 +1,26 @@
+from app.rag.tools import CALCULATOR_TOOL, calculate_expression, execute_tool
+
+
+def test_calculator_tool_evaluates_basic_expression():
+    assert calculate_expression("1000 - 250") == "750"
+    assert calculate_expression("10 + 5 * 2") == "20"
+    assert calculate_expression("10 / 4") == "2.5"
+
+
+def test_calculator_tool_rejects_unsafe_expression():
+    try:
+        calculate_expression("__import__('os').system('echo x')")
+    except ValueError as exc:
+        assert "Invalid calculator expression" in str(exc) or "Unsupported expression" in str(exc)
+    else:
+        assert False, "Unsafe expressions should not be evaluated"
+
+
+def test_execute_tool_with_expression_argument():
+    result = execute_tool("calculator", {"expression": "12 * 3"})
+    assert result == "36"
+
+
+def test_calculator_tool_metadata():
+    assert CALCULATOR_TOOL["function"]["name"] == "calculator"
+    assert "expression" in CALCULATOR_TOOL["function"]["parameters"]["properties"]