File size: 4,374 Bytes
f2eba97
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
129
130
import re
import ast
import operator
from typing import List, Dict
from groq import Groq
from agents.base_agent import BaseAgent
import config
from utils.logger import get_logger

logger = get_logger(__name__)

_SAFE_OPERATORS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Pow: operator.pow,
    ast.USub: operator.neg,
    ast.Mod: operator.mod,
    ast.FloorDiv: operator.floordiv,
}

_PURE_MATH_RE = re.compile(
    r"^\s*[\d\s\.\+\-\*\/\(\)\^%]+\s*$"
)

def _safe_eval(expr: str) -> float | None:
    expr = expr.replace("^", "**")

    def _eval_node(node):
        if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
            return node.value
        elif isinstance(node, ast.BinOp):
            op_func = _SAFE_OPERATORS.get(type(node.op))
            if op_func is None:
                raise ValueError(f"Unsupported operator: {type(node.op)}")
            return op_func(_eval_node(node.left), _eval_node(node.right))
        elif isinstance(node, ast.UnaryOp):
            op_func = _SAFE_OPERATORS.get(type(node.op))
            if op_func is None:
                raise ValueError(f"Unsupported unary operator: {type(node.op)}")
            return op_func(_eval_node(node.operand))
        else:
            raise ValueError(f"Unsupported AST node: {type(node)}")

    try:
        tree = ast.parse(expr.strip(), mode="eval")
        result = _eval_node(tree.body)
        return result
    except Exception:
        return None

def _extract_expression(query: str) -> str | None:
    cleaned = re.sub(
        r"(?i)^(what\s+is|calculate|compute|evaluate|solve|find)\s*:?\s*", "", query.strip()
    )
    cleaned = re.sub(r"[?!]+$", "", cleaned).strip()

    if _PURE_MATH_RE.match(cleaned):
        return cleaned

    match = re.search(r"[\d]+\s*[\+\-\*\/\^%]\s*[\d\.\s\+\-\*\/\^\(\)%]+", query)
    if match:
        return match.group(0).strip()

    return None

class MathAgent(BaseAgent):
    SYSTEM_PROMPT = """You are a precise mathematical reasoning assistant.
Solve the given problem step by step using chain-of-thought reasoning.
Always:
  1. Identify the mathematical concept involved
  2. Show your working step by step
  3. State your final answer clearly on a new line starting with "Answer:"
Be concise but thorough. Always include a space after bolding (e.g., **Calculation:** result) for correct UI rendering."""

    def __init__(self):
        super().__init__(name="Math Agent")
        self._client = Groq(api_key=config.GROQ_API_KEY)
        logger.info("MathAgent ready.")

    def run(
        self,
        query: str,
        context: str = "",
        history: List[Dict[str, str]] = None,
        session_id: str = "",
    ) -> str:
        logger.info(f"MathAgent processing: '{query[:80]}'")

        expr = _extract_expression(query)
        if expr is not None:
            result = _safe_eval(expr)
            if result is not None:
                formatted = int(result) if result == int(result) else round(result, 6)
                answer = (
                    f"**Calculation:** `{expr.strip()} = {formatted}`\n\n"
                    f"**Answer: {formatted}**\n\n"
                    f"*(Computed via Python arithmetic — 100% accurate)*"
                )
                logger.info(f"  → Computed via Python eval: {formatted}")
                return answer

        logger.info("  → Falling back to LLM chain-of-thought reasoning.")

        messages = [{"role": "system", "content": self.SYSTEM_PROMPT}]

        if history:
            messages.extend(history[-4:])

        messages.append({"role": "user", "content": query})

        try:
            response = self._client.chat.completions.create(
                model=config.GROQ_MODEL_NAME,
                messages=messages,
                temperature=0.1,
                max_tokens=config.GROQ_MAX_TOKENS,
            )
            answer = response.choices[0].message.content.strip()
            logger.info("  → LLM math response received.")
            return answer

        except Exception as exc:
            logger.error(f"MathAgent LLM call failed: {exc}")
            return (
                f"I encountered an error processing your math query. "
                f"Please try rephrasing it.\n\nError: {exc}"
            )