Spaces:
Build error
Build error
| # chatbot_agent.py | |
| import os | |
| import json | |
| import re | |
| from openai import OpenAI | |
| import traceback | |
| import logging | |
| # Get logger | |
| logger = logging.getLogger(__name__) | |
| client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) | |
| def format_message(role, content): | |
| """Format message for chatbot history.""" | |
| return {"role": role, "content": content} | |
| def initialize_graph_prompt(graph_data): | |
| """Initialize the conversation with detailed node, edge, linker, and text information.""" | |
| summary_info = ( | |
| f"Symbols (represented as nodes): {graph_data['summary']['symbol_count']}, " | |
| f"Texts: {graph_data['summary']['text_count']}, " | |
| f"Lines: {graph_data['summary']['line_count']}, " | |
| f"Linkers: {graph_data['summary']['linker_count']}, " | |
| f"Edges: {graph_data['summary']['edge_count']}." | |
| ) | |
| # Prepare detailed node (symbol) data | |
| node_details = "Nodes (symbols) in the graph include the following details:\n" | |
| for symbol in graph_data["detailed_results"]["symbols"]: | |
| node_details += ( | |
| f"Node ID: {symbol['symbol_id']}, Class ID: {symbol['class_id']}, " | |
| f"Category: {symbol['category']}, Type: {symbol['type']}, " | |
| f"Label: {symbol['label']}, Confidence: {symbol['confidence']}\n" | |
| ) | |
| # Prepare edge data | |
| edge_details = "Edges in the graph showing connections between nodes are as follows:\n" | |
| for edge in graph_data["detailed_results"].get("edges", []): | |
| edge_details += ( | |
| f"Edge ID: {edge['edge_id']}, From Node: {edge['symbol_1_id']}, " | |
| f"To Node: {edge['symbol_2_id']}, Type: {edge.get('type', 'unknown')}\n" | |
| ) | |
| # Prepare linker data | |
| linker_details = "Linkers in the diagram are as follows:\n" | |
| for linker in graph_data["detailed_results"].get("linkers", []): | |
| linker_details += ( | |
| f"Symbol ID: {linker['symbol_id']}, Associated Text IDs: {linker.get('text_ids', [])}, " | |
| f"Associated Edge IDs: {linker.get('edge_ids', [])}, Position: {linker.get('bbox', 'unknown')}\n" | |
| ) | |
| # Prepare text (tag) data | |
| text_details = "Text elements with associated tags in the diagram are as follows:\n" | |
| for text in graph_data["detailed_results"].get("texts", []): | |
| text_details += ( | |
| f"Text ID: {text['text_id']}, Content: {text['content']}, " | |
| f"Confidence: {text['confidence']}, Position: {text['bbox']}\n" | |
| ) | |
| initial_prompt = ( | |
| "You have access to a knowledge graph generated from a P&ID diagram. " | |
| f"The summary information includes:\n{summary_info}\n\n" | |
| "The detailed information about each node (symbol) in the graph is as follows:\n" | |
| f"{node_details}\n" | |
| "The edges connecting these nodes are as follows:\n" | |
| f"{edge_details}\n" | |
| "The linkers in the diagram are as follows:\n" | |
| f"{linker_details}\n" | |
| "The text elements and their tags in the diagram are as follows:\n" | |
| f"{text_details}\n" | |
| "Answer questions about specific nodes, edges, types, labels, categories, linkers, or text tags using this information." | |
| ) | |
| return initial_prompt | |
| def get_assistant_response(user_message, json_path): | |
| """Generate response based on P&ID data and OpenAI.""" | |
| try: | |
| # Load the aggregated data | |
| with open(json_path, 'r') as f: | |
| data = json.load(f) | |
| # Process the user's question | |
| question = user_message.lower() | |
| # Use rule-based responses for specific questions | |
| if "valve" in question or "valves" in question: | |
| valve_count = sum(1 for symbol in data.get('symbols', []) | |
| if 'class' in symbol and 'valve' in symbol['class'].lower()) | |
| return f"I found {valve_count} valves in this P&ID." | |
| elif "pump" in question or "pumps" in question: | |
| pump_count = sum(1 for symbol in data.get('symbols', []) | |
| if 'class' in symbol and 'pump' in symbol['class'].lower()) | |
| return f"I found {pump_count} pumps in this P&ID." | |
| elif "equipment" in question or "components" in question: | |
| equipment_types = {} | |
| for symbol in data.get('symbols', []): | |
| if 'class' in symbol: | |
| eq_type = symbol['class'] | |
| equipment_types[eq_type] = equipment_types.get(eq_type, 0) + 1 | |
| response = "Here's a summary of the equipment I found:\n" | |
| for eq_type, count in equipment_types.items(): | |
| response += f"- {eq_type}: {count}\n" | |
| return response | |
| # For other questions, use OpenAI | |
| else: | |
| # Prepare the conversation context | |
| graph_data = { | |
| "summary": { | |
| "symbol_count": len(data.get('symbols', [])), | |
| "text_count": len(data.get('texts', [])), | |
| "line_count": len(data.get('lines', [])), | |
| "edge_count": len(data.get('edges', [])), | |
| }, | |
| "detailed_results": data | |
| } | |
| initial_prompt = initialize_graph_prompt(graph_data) | |
| conversation = [ | |
| {"role": "system", "content": initial_prompt}, | |
| {"role": "user", "content": user_message} | |
| ] | |
| response = client.chat.completions.create( | |
| model="gpt-4-turbo", | |
| messages=conversation | |
| ) | |
| return response.choices[0].message.content | |
| except Exception as e: | |
| logger.error(f"Error in get_assistant_response: {str(e)}") | |
| logger.error(traceback.format_exc()) | |
| return "I apologize, but I encountered an error analyzing the P&ID data. Please try asking a different question." | |
| # Testing and Usage block | |
| if __name__ == "__main__": | |
| # Load the knowledge graph data from JSON file | |
| json_file_path = "results/0_aggregated_detections.json" | |
| try: | |
| with open(json_file_path, 'r') as file: | |
| graph_data = json.load(file) | |
| except FileNotFoundError: | |
| print(f"Error: File not found at {json_file_path}") | |
| graph_data = None | |
| except json.JSONDecodeError: | |
| print("Error: Failed to decode JSON. Please check the file format.") | |
| graph_data = None | |
| # Initialize conversation history with assistant's welcome message | |
| 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.")] | |
| # Print the assistant's welcome message | |
| print("Assistant:", history[0]["content"]) | |
| # Individual Testing Options | |
| if graph_data: | |
| # Option 1: Test the graph prompt initialization | |
| print("\n--- Test: Graph Prompt Initialization ---") | |
| initial_prompt = initialize_graph_prompt(graph_data) | |
| print(initial_prompt) | |
| # Option 2: Simulate a conversation with a test question | |
| print("\n--- Test: Simulate Conversation ---") | |
| test_question = "Can you tell me about the connections between the nodes?" | |
| history.append(format_message("user", test_question)) | |
| print(f"\nUser: {test_question}") | |
| for response in get_assistant_response(test_question, json_file_path): | |
| print("Assistant:", response) | |
| history.append(format_message("assistant", response)) | |
| # Option 3: Manually input questions for interactive testing | |
| while True: | |
| user_question = input("\nYou: ") | |
| if user_question.lower() in ["exit", "quit"]: | |
| print("Exiting chat. Goodbye!") | |
| break | |
| history.append(format_message("user", user_question)) | |
| for response in get_assistant_response(user_question, json_file_path): | |
| print("Assistant:", response) | |
| history.append(format_message("assistant", response)) | |
| else: | |
| print("Unable to load graph data. Please check the file path and format.") | |