File size: 7,155 Bytes
790e0e9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
"""
AI Agent - OpenAI-powered assistant with function calling
"""
import json
from typing import List, Dict, Any, Optional
import logging
from openai import OpenAI

from config import OPENAI_API_KEY, OPENAI_MODEL
from agent.tools import AgentTools


class AIAgent:
    """AI Agent powered by OpenAI with function calling capabilities"""
    
    SYSTEM_PROMPT = """You are a helpful data analyst assistant for a car auction/pricing database. 
Your role is to help users understand and query car pricing data.

IMPORTANT GUIDELINES:
1. **Data Privacy**: Never pass the entire dataset to your responses. Only use the tools to query specific data.
2. **Safety**: You can only execute SELECT queries. Any attempt to modify data (DELETE, UPDATE, INSERT, DROP) will be blocked.
3. **Tool Usage**: 
   - Use `query_database` for specific data queries
   - Use `get_database_statistics` for general overviews and statistics
   - Use `generate_chart` when the user asks for a chart, visualization, or trend analysis. Choose the most appropriate chart type (bar, column, line, pie, scatter).
   - Use `create_support_ticket` when you cannot help or user requests human assistance
4. **Support Escalation**: If you cannot answer a question or the user seems frustrated, proactively suggest creating a support ticket.
5. **Clear Communication**: Explain your findings clearly with relevant numbers and insights.

DATABASE SCHEMA:
- Table: cars
- Columns: year, make, model, trim, body, transmission, vin, state, condition, odometer, color, interior, seller, mmr, sellingprice, saledate

Be concise, helpful, and data-driven in your responses."""
    
    def __init__(self, tools: AgentTools):
        self.tools = tools
        self.client = OpenAI(api_key=OPENAI_API_KEY)
        self.model = OPENAI_MODEL
        self.logger = logging.getLogger(__name__)
        self.conversation_history: List[Dict[str, Any]] = []
        
        # Initialize with system prompt
        self.conversation_history.append({
            "role": "system",
            "content": self.SYSTEM_PROMPT
        })
    
    def chat(self, user_message: str) -> Dict[str, Any]:
        """
        Process a user message and return AI response with metadata
        
        Args:
            user_message: User's question or request
            
        Returns:
            Dictionary with 'content' (str) and optional 'chart' (dict)
        """
        try:
            # Add user message to history
            self.conversation_history.append({
                "role": "user",
                "content": user_message
            })
            
            # Get AI response with function calling
            return self._get_ai_response()
            
        except Exception as e:
            error_msg = f"Error processing message: {str(e)}"
            self.logger.error(error_msg)
            return {
                "content": f"❌ {error_msg}",
                "chart": None
            }
    
    def _get_ai_response(self, max_iterations: int = 5) -> Dict[str, Any]:
        """
        Get AI response with function calling loop
        
        Args:
            max_iterations: Maximum number of function calling iterations
            
        Returns:
            Dictionary with 'content' and optional 'chart'
        """
        iteration = 0
        last_chart = None
        
        while iteration < max_iterations:
            iteration += 1
            
            # Call OpenAI API
            response = self.client.chat.completions.create(
                model=self.model,
                messages=self.conversation_history,
                tools=AgentTools.get_tool_definitions(),
                tool_choice="auto"
            )
            
            message = response.choices[0].message
            
            # Check if AI wants to call a function
            if message.tool_calls:
                # Add assistant message to history
                self.conversation_history.append({
                    "role": "assistant",
                    "content": message.content,
                    "tool_calls": [
                        {
                            "id": tc.id,
                            "type": tc.type,
                            "function": {
                                "name": tc.function.name,
                                "arguments": tc.function.arguments
                            }
                        }
                        for tc in message.tool_calls
                    ]
                })
                
                # Execute each tool call
                for tool_call in message.tool_calls:
                    function_name = tool_call.function.name
                    function_args = json.loads(tool_call.function.arguments)
                    
                    self.logger.info(f"AI calling function: {function_name}")
                    
                    # Execute the tool
                    result = self.tools.execute_tool(function_name, function_args)
                    
                    # Capture chart result if it's a chart
                    if result.get('is_chart'):
                        last_chart = result.get('chart_config')
                    
                    # Add function result to history
                    self.conversation_history.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": json.dumps(result)
                    })
                
                # Continue loop to get final response
                continue
            
            else:
                # No more function calls, return final response
                final_response = message.content or "I apologize, but I couldn't generate a response."
                
                # Add to history
                self.conversation_history.append({
                    "role": "assistant",
                    "content": final_response
                })
                
                return {
                    "content": final_response,
                    "chart": last_chart
                }
        
        # Max iterations reached
        return {
            "content": "I apologize, but I'm having trouble processing your request. Would you like me to create a support ticket for human assistance?",
            "chart": None
        }
    
    def reset_conversation(self):
        """Reset conversation history"""
        self.conversation_history = [{
            "role": "system",
            "content": self.SYSTEM_PROMPT
        }]
        self.logger.info("Conversation history reset")
    
    def get_conversation_context(self) -> str:
        """Get conversation history as formatted string for support tickets"""
        context = []
        for msg in self.conversation_history:
            if msg["role"] == "user":
                context.append(f"User: {msg['content']}")
            elif msg["role"] == "assistant" and msg.get("content"):
                context.append(f"Assistant: {msg['content']}")
        
        return "\n\n".join(context)