File size: 3,264 Bytes
e77665f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
131
132
133
134
135
136
137
138
139
140
141
import os
from urllib import response
from dotenv import load_dotenv
from openai import OpenAI
from backend.app.tools.tool_registry import ToolRegistry
from configs.logging_config import setup_logger
from configs.settings import config

load_dotenv()

logger = setup_logger("tool_agent")

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

MODEL_NAME = config["llm"]["model"]


class ToolAgent:
    def __init__(self):
        self.registry = ToolRegistry()
    def decide_tool_llm(self, query, llm_model=MODEL_NAME):
        tools = [
            "analyze_revenue_trend",
            "compare_regions",
            "detect_anomalies"
        ]

        prompt = f"""
    You are an AI routing agent.

    Available tools:
    - analyze_revenue_trend: Detect revenue increase or decrease
    - compare_regions: Compare performance across regions
    - detect_anomalies: Find unusual revenue patterns

    User Query:
    {query}

    Return ONLY one tool name from the list above.
    If no tool is relevant, return "NONE".
    """

        response = client.chat.completions.create(
            model=llm_model,
            messages=[{"role": "user", "content": prompt}]
        )

        tool_name = response.choices[0].message.content.strip()

        # ✅ Normalize output
        tool_name = tool_name.replace('"', '').strip()

        # ✅ Validate tool
        if tool_name not in tools:
            logger.warning(f"Invalid tool from LLM: {tool_name}")
            return None

        return tool_name

    def decide_tool(self, query):
        query = query.lower()

        keywords = {
    "analyze_revenue_trend": ["trend", "drop", "increase"],
    "compare_regions": ["region", "area"],
    "detect_anomalies": ["anomaly", "anomalies", "outlier", "unusual"]
}

        for tool, words in keywords.items():
            if any(word in query for word in words):
                return tool

    def run(self, query):
        logger.info(f"User Query: {query}")

        tool_name = self.decide_tool_llm(query)

        logger.info(f"Selected Tool: {tool_name}")

        if tool_name is None:
            logger.warning("No tool matched. Falling back to LLM.")

        if tool_name:
            tool = self.registry.get_tool(tool_name)
            tool_result = tool()

            logger.info(f"Tool Result: {tool_result}")

            prompt = f"""
You are an AI Operations Copilot.

User Question:
{query}

Tool Output:
{tool_result}

Explain clearly using the tool output.
Provide:
- Summary
- Key Insight
- Suggested Action
"""

        else:
            prompt = f"""
User Question:
{query}

Respond normally.
"""

        response = client.chat.completions.create(
            model="gpt-4.1-mini",
            messages=[{"role": "user", "content": prompt}]
        )

        output = response.choices[0].message.content
        logger.info(f"Response: {output}")

        return output


def run_cli():
    agent = ToolAgent()

    print("=== Tool Agent ===")
    print("Type 'exit' to quit\n")

    while True:
        query = input("Enter your query: ")

        if query.lower() == "exit":
            break

        result = agent.run(query)
        print(f"\nAgent Response:\n{result}\n")


if __name__ == "__main__":
    run_cli()