Decipher/generator.py

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