src/tig_engine.py

import os
import json
import re
from dotenv import load_dotenv
from google import genai
from openai import OpenAI

class IntelliMod:
    def __init__(self):
        load_dotenv()
        self.gemini_key = os.getenv("GEMINI_API_KEY")
        self.openrouter_key = os.getenv("OPENROUTER_API_KEY")
        
        self.gemini_client = None
        if self.gemini_key:
            self.gemini_client = genai.Client(api_key=self.gemini_key)
            
        self.openrouter_client = None
        if self.openrouter_key:
            self.openrouter_client = OpenAI(
                api_key=self.openrouter_key,
                base_url="https://openrouter.ai/api/v1"
            )

        self.tool_registry = {
            "coding": "anthropic/claude-sonnet-4",
            "planning": "anthropic/claude-opus-4",
            "creative_text": "openai/gpt-5.1-chat",
            "research": "gemini-3-flash-preview",
            "chat": "gemini-2.5-flash"
        }
        
        self.active_model = self.tool_registry["chat"]
        # Track conversation state for the prompt builder flow
        self.prompt_state = {}

    def _call_model(self, model, system_prompt, user_prompt):
        """Core execution — routes to the right model with optional system prompt."""
        # PATH A: Google Direct
        if "gemini" in model and self.gemini_client:
            try:
                combined = f"{system_prompt}\n\n{user_prompt}" if system_prompt else user_prompt
                response = self.gemini_client.models.generate_content(
                    model=model,
                    contents=combined
                )
                return response.text
            except Exception as e:
                print(f"  [IntelliMod] Google ({model}) failed: {e}")

        # PATH B: OpenRouter
        if self.openrouter_client:
            try:
                messages = []
                if system_prompt:
                    messages.append({"role": "system", "content": system_prompt})
                messages.append({"role": "user", "content": user_prompt})
                response = self.openrouter_client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=0.7
                )
                return response.choices[0].message.content
            except Exception as e:
                return f"[Error] OpenRouter failed: {e}"

        return "[Error] No API clients configured."

    # ──────────────────────────────────────────────
    # PROMPT BUILDER — Interactive prompt crafting
    # ──────────────────────────────────────────────

    PROMPT_BUILDER_SYSTEM = """You are IntelliMod, a prompt engineering assistant. Your job is to help the user craft effective prompts for AI models.

ADAPTIVE BEHAVIOR:
- If the user gives a clear, detailed request with enough context — skip questions and build the prompt immediately.
- If the request is vague, short, or missing critical details — ask questions ONE AT A TIME to fill gaps.

Questions to consider (only if needed):
- Goal: What should the output accomplish?
- Audience: Who is the output for?
- Format: Text, code, JSON, table, list?
- Tone: Formal, casual, technical, creative?
- Constraints: Length limits, exclusions, specific requirements?

Ask ONE question at a time. Don't overwhelm the user.
When you have enough info to build, say "READY TO BUILD" on its own line and then output the crafted prompt with clear structure.
"""

    def run_prompt_builder(self, user_input, session_id="default"):
        """Manages the interactive prompt builder conversation."""
        if session_id not in self.prompt_state:
            self.prompt_state[session_id] = {"history": [], "ready": False, "compiled": None}
        
        state = self.prompt_state[session_id]
        state["history"].append({"role": "user", "content": user_input})
        
        # Build context from history
        context = "\n".join([
            f"{'User' if m['role']=='user' else 'You'}: {m['content']}"
            for m in state["history"]
        ])

        full_prompt = f"""Previous conversation:
{context}

{'The user has answered all questions. Say "READY TO BUILD" and output the compiled prompt.' if len(state['history']) >= 4 else 'Continue the conversation. Ask the next question if needed.'}"""

        response = self._call_model("gemini-2.5-flash", self.PROMPT_BUILDER_SYSTEM, full_prompt)
        
        # Check if it's ready to build
        if "READY TO BUILD" in response.upper():
            state["ready"] = True
            # Extract the prompt after "READY TO BUILD"
            parts = re.split(r'(?i)READY TO BUILD', response, maxsplit=1)
            state["compiled"] = parts[1].strip() if len(parts) > 1 else response
        
        state["history"].append({"role": "assistant", "content": response})
        return response, state["ready"], state.get("compiled")

    # ──────────────────────────────────────────────
    # TIG PIPELINE (Legacy — direct Q&A mode)
    # ──────────────────────────────────────────────

    def detect_intent(self, user_prompt):
        if len(user_prompt.split()) < 5:
            return "chat"
        if not self.gemini_client:
            return "chat"
        try:
            classifier_prompt = f"""Classify this prompt into ONE word: [coding, creative_text, research, planning, chat]

- coding: python, scripts, html, debugging, logic, "write code"
- creative_text: stories, essays, poems, writing
- research: facts, history, summarizing, looking up info
- planning: complex step-by-step plans, architecture
- chat: casual conversation, greetings, simple questions, thank yous

PROMPT: "{user_prompt[:500]}"
Return only the classification word."""
            response = self.gemini_client.models.generate_content(
                model="gemini-2.5-flash",
                contents=classifier_prompt
            )
            intent = response.text.strip().lower()
            for valid in self.tool_registry.keys():
                if valid in intent:
                    return valid
            return "chat"
        except Exception as e:
            print(f"  [IntelliMod] Intent detection failed: {e}")
            return "chat"

    def run_tig_pipeline(self, prompt, force_model=None):
        """Direct Q&A mode — sends prompt to model, returns raw response."""
        if force_model:
            target_model = force_model
        else:
            intent = self.detect_intent(prompt)
            target_model = self.tool_registry.get(intent, "gemini-2.5-flash")
        
        self.active_model = target_model
        
        # System prompt to make it a useful assistant
        system = "You are a helpful AI assistant. Answer concisely and accurately."
        return self._call_model(target_model, system, prompt)

    # ──────────────────────────────────────────────
    # MPA AUDITOR
    # ──────────────────────────────────────────────

    MPA_PATH = os.path.join(os.path.dirname(__file__), "..", "intellimod_system", "content", "intellimod_knowledge", "audit-protocols", "modular-prompt-auditor.md")

    def run_mpa_pipeline(self, user_prompt, force_model=None):
        try:
            with open(self.MPA_PATH, "r") as f:
                mpa_spec = f.read()
        except Exception as e:
            return f"[MPA Error] Could not load auditor spec: {e}"

        cutoff = mpa_spec.find("Insert Prompt to Evaluate")
        if cutoff != -1:
            mpa_spec = mpa_spec[:cutoff]

        evaluator_prompt = f"""{mpa_spec.strip()}

---
## EVALUATION TARGET

Prompt to Evaluate:
```
{user_prompt}
```

---
## INSTRUCTIONS

Execute ALL 9 steps above against this prompt.
Be thorough and specific. Output each step clearly.
"""
        target_model = force_model or "anthropic/claude-sonnet-4"
        return self._call_model(target_model, "", evaluator_prompt)