msintui/Intelligent_PID
0
1# chatbot_agent.py2 3import os4import json5import re6from openai import OpenAI7import traceback8import logging9from dotenv import load_dotenv10 11# Load environment variables12load_dotenv()13 14# Get logger15logger = logging.getLogger(__name__)16 17# Initialize OpenAI client with error handling18def get_openai_client():19 api_key = os.getenv("OPENAI_API_KEY")20 if not api_key:21 raise ValueError("OpenAI API key not found in environment variables")22 return OpenAI(api_key=api_key)23 24def format_message(role, content):25 """Format message for chatbot history."""26 return {"role": role, "content": content}27 28def initialize_graph_prompt(graph_data):29 """Initialize the conversation with available node and edge information."""30 try:31 # Get summary info with safe fallbacks32 summary = graph_data.get('summary', {})33 summary_parts = []34 35 # Only include counts that exist36 if 'symbol_count' in summary:37 summary_parts.append(f"Symbols: {summary['symbol_count']}")38 if 'text_count' in summary:39 summary_parts.append(f"Texts: {summary['text_count']}")40 if 'line_count' in summary:41 summary_parts.append(f"Lines: {summary['line_count']}")42 if 'edge_count' in summary:43 summary_parts.append(f"Edges: {summary['edge_count']}")44 45 summary_info = ", ".join(summary_parts) + "."46 47 # Prepare node details only if they exist48 node_details = ""49 detailed_results = graph_data.get('detailed_results', {})50 if 'symbols' in detailed_results:51 node_details = "Nodes (symbols) in the graph include:\n"52 for symbol in detailed_results['symbols']:53 details = []54 if 'symbol_id' in symbol:55 details.append(f"ID: {symbol['symbol_id']}")56 if 'class_id' in symbol:57 details.append(f"Class: {symbol['class_id']}")58 if 'category' in symbol:59 details.append(f"Category: {symbol['category']}")60 if 'type' in symbol:61 details.append(f"Type: {symbol['type']}")62 if 'label' in symbol:63 details.append(f"Label: {symbol['label']}")64 if details: # Only add if we have any details65 node_details += ", ".join(details) + "\n"66 67 initial_prompt = (68 "You have access to a knowledge graph generated from a P&ID diagram. "69 f"The summary information includes:\n{summary_info}\n\n"70 f"{node_details}\n"71 "Answer questions about the P&ID elements using this information."72 )73 74 return initial_prompt75 76 except Exception as e:77 logger.error(f"Error creating initial prompt: {str(e)}")78 return ("I have access to a P&ID diagram knowledge graph. "79 "I can help answer questions about the diagram elements.")80 81def get_assistant_response(user_message, json_path):82 """Generate response based on P&ID data and OpenAI."""83 try:84 client = get_openai_client()85 # Load the aggregated data86 with open(json_path, 'r') as f:87 data = json.load(f)88 89 # Process the user's question90 question = user_message.lower()91 92 # Use rule-based responses for specific questions93 if "valve" in question or "valves" in question:94 valve_count = sum(1 for symbol in data.get('symbols', []) 95 if 'class' in symbol and 'valve' in symbol['class'].lower())96 return f"I found {valve_count} valves in this P&ID."97 98 elif "pump" in question or "pumps" in question:99 pump_count = sum(1 for symbol in data.get('symbols', [])100 if 'class' in symbol and 'pump' in symbol['class'].lower())101 return f"I found {pump_count} pumps in this P&ID."102 103 elif "equipment" in question or "components" in question:104 equipment_types = {}105 for symbol in data.get('symbols', []):106 if 'class' in symbol:107 eq_type = symbol['class']108 equipment_types[eq_type] = equipment_types.get(eq_type, 0) + 1109 110 response = "Here's a summary of the equipment I found:\n"111 for eq_type, count in equipment_types.items():112 response += f"- {eq_type}: {count}\n"113 return response114 115 # For other questions, use OpenAI116 else:117 # Prepare the conversation context118 graph_data = {119 "summary": {120 "symbol_count": len(data.get('symbols', [])),121 "text_count": len(data.get('texts', [])),122 "line_count": len(data.get('lines', [])),123 "edge_count": len(data.get('edges', [])),124 },125 "detailed_results": data126 }127 128 initial_prompt = initialize_graph_prompt(graph_data)129 conversation = [130 {"role": "system", "content": initial_prompt},131 {"role": "user", "content": user_message}132 ]133 134 response = client.chat.completions.create(135 model="gpt-4-turbo",136 messages=conversation137 )138 return response.choices[0].message.content139 140 except Exception as e:141 logger.error(f"Error in get_assistant_response: {str(e)}")142 logger.error(traceback.format_exc())143 return "I apologize, but I encountered an error analyzing the P&ID data. Please try asking a different question."144 145# Testing and Usage block146if __name__ == "__main__":147 # Load the knowledge graph data from JSON file148 json_file_path = "results/0_aggregated_detections.json"149 try:150 with open(json_file_path, 'r') as file:151 graph_data = json.load(file)152 except FileNotFoundError:153 print(f"Error: File not found at {json_file_path}")154 graph_data = None155 except json.JSONDecodeError:156 print("Error: Failed to decode JSON. Please check the file format.")157 graph_data = None158 159 # Initialize conversation history with assistant's welcome message160 history = [format_message("assistant", "Hello! I am ready to answer your questions about the P&ID knowledge graph. The graph includes nodes (symbols), edges, linkers, and text tags, and I have detailed information available about each. Please ask any questions related to these elements and their connections.")]161 162 # Print the assistant's welcome message163 print("Assistant:", history[0]["content"])164 165 # Individual Testing Options166 if graph_data:167 # Option 1: Test the graph prompt initialization168 print("\n--- Test: Graph Prompt Initialization ---")169 initial_prompt = initialize_graph_prompt(graph_data)170 print(initial_prompt)171 172 # Option 2: Simulate a conversation with a test question173 print("\n--- Test: Simulate Conversation ---")174 test_question = "Can you tell me about the connections between the nodes?"175 history.append(format_message("user", test_question))176 177 print(f"\nUser: {test_question}")178 for response in get_assistant_response(test_question, json_file_path):179 print("Assistant:", response)180 history.append(format_message("assistant", response))181 182 # Option 3: Manually input questions for interactive testing183 while True:184 user_question = input("\nYou: ")185 if user_question.lower() in ["exit", "quit"]:186 print("Exiting chat. Goodbye!")187 break188 189 history.append(format_message("user", user_question))190 for response in get_assistant_response(user_question, json_file_path):191 print("Assistant:", response)192 history.append(format_message("assistant", response))193 else:194 print("Unable to load graph data. Please check the file path and format.")195