import os import json from openai import OpenAI from dotenv import load_dotenv from Decipher.models import NetworkTopology from Decipher.validator import validate_topology load_dotenv() client = OpenAI( api_key=os.getenv("GHAYMAH_GENERATE_AI_KEY"), base_url=os.getenv("GHAYMAH_BASE_URL"), timeout=60.0 # Increase timeout to 60 seconds ) SYSTEM_PROMPT = """ You are a Senior Network Architect. Your task is to generate a highly accurate and valid network topology JSON based on a user's description. The output MUST strictly follow the provided JSON schema. Principles for accuracy: 1. IP Addressing: Ensure devices in the same segment have unique IPs in the same subnet. 2. Connectivity: Define physical links correctly between existing devices and ports. 3. Routing: If OSPF or static routes are requested, ensure the configuration is logically sound. 4. VLANs: Ensure ports are assigned to correct VLANs if specified. 5. Device Types: Use realistic device types (Router, Switch, PC, Server, etc.) and models (e.g., 2911, 2960). You must respond ONLY with the valid JSON object. No markdown, no conversational text. Ensure all relevant arrays (devices, physical_links) are populated based on the user's request. Do not return empty arrays if the user asked for specific components. Example Output Structure: { "version": "1.0", "devices": [ { "name": "Router1", "type": "Router", "model": "2911", "interfaces": [{"name": "GigabitEthernet0/0", "ip": "192.168.1.1", "subnet": "255.255.255.0"}], "routing": {"static_routes": [], "ospf": [{"process_id": "1", "networks": ["192.168.1.0 0.0.0.255 area 0"]}]} } ], "physical_links": [ {"type": "Copper", "from_device": "Router1", "to_device": "Switch1", "from_port": "GigabitEthernet0/0", "to_port": "FastEthernet0/1"} ], "vlans": {"10": "Sales", "20": "HR"}, "instructions": "Brief summary of the network..." } """ def generate_network(prompt: str) -> NetworkTopology: try: response = client.chat.completions.create( model="DeepSeek-V3-0324", # Using available model messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": f"Generate a network topology for: {prompt}"} ], response_format={"type": "json_object"} ) content = response.choices[0].message.content print(f"--- RAW LLM RESPONSE ---\n{content}\n-----------------------") data = json.loads(content) if not data.get("devices") and not data.get("physical_links"): # If AI returned an empty-ish object, it might have failed to understand. # We can try to raise an error so the API doesn't return an empty success. raise ValueError("AI generated an empty network topology. Please try a more descriptive prompt.") # Validate and parse into Pydantic model topology = NetworkTopology(**data) # Intelligent Validation Layer validation_issues = validate_topology(topology) if validation_issues: # Append issues to instructions so user is aware of logical warnings issue_text = "\n\n[Accuracy Layer Warnings]:\n- " + "\n- ".join(validation_issues) topology.instructions += issue_text return topology except Exception as e: print(f"Error generating network: {e}") raise e