agent.py

import os
import json
from openai import OpenAI
from dotenv import load_dotenv
import piko
import prompt_router
load_dotenv()

# Configure OpenAI client for Nvidia's API
nvidia_key = os.getenv("NVIDIA_API_KEY") or os.getenv("API_KEY")
piko_key = os.getenv("PIKO_API_KEY")

if not nvidia_key:
    print("⚠️  Warning: NVIDIA_API_KEY or API_KEY not found in environment.")
    print("If you are using OpenAI directly, please set OPENAI_API_KEY.")
    nvidia_key = os.getenv("OPENAI_API_KEY")

if not nvidia_key:
    print("❌ Critical: No LLM API Key found. Please update your .env file.")
    # We will let the OpenAI client raise the error if still missing, or handle it here
    
client = OpenAI(
  base_url="https://integrate.api.nvidia.com/v1" if os.getenv("NVIDIA_API_KEY") else "https://api.openai.com/v1",
  api_key=nvidia_key,
  timeout=120.0 # Increased timeout for long-reasoning models
)

# Note: PIKO_API_KEY will be provided dynamically per session.

# --- Optimized Model Parameters ---
# Temperature: 0.2 for deterministic, structured JSON output
# Max tokens: 800 — enough for any action JSON. Model must output JSON first, no preamble.
MAIN_MODEL_TEMP = 0.2
MAIN_MODEL_TOP_P = 0.9
MAIN_MODEL_MAX_TOKENS = 800


def _load_base_prompt():
    """Load the raw base system prompt from file."""
    base_dir = os.path.dirname(os.path.abspath(__file__))
    prompt = "You are Piko, an AI backend automation assistant for Postpipe."
    
    prompt_path = os.path.join(base_dir, "piko_system_prompt.txt")
    try:
        if os.path.exists(prompt_path):
            with open(prompt_path, "r", encoding="utf-8") as f:
                prompt = f.read()
    except Exception as e:
        print(f"Error loading system prompt: {e}")
        
    return prompt

# Cache the base prompt at module load
_BASE_PROMPT = _load_base_prompt()

def get_system_prompt():
    return _BASE_PROMPT

def build_contextual_prompt(user_message: str) -> str:
    """
    Divide-and-conquer prompt builder.
    Uses the prompt router to analyze the user message and inject
    only the relevant ecommerce chunks into the system prompt.
    """
    return prompt_router.build_system_prompt(_BASE_PROMPT, user_message)


def extract_action(text):
    import re
    # 1. Look for markdown json blocks first
    json_block_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
    if json_block_match:
        try:
            return json.loads(json_block_match.group(1))
        except:
            pass
            
    # 2. Fallback: try to find anything that looks like a JSON block
    json_matches = list(re.finditer(r'\{.*\}', text, re.DOTALL))
    if json_matches:
        try:
            return json.loads(json_matches[-1].group())
        except:
            pass
            
    # 2. Fallback: Check for XML-style tool calls (StepFun format)
    # Format: <function=action_name>params_as_json_or_empty</function>
    xml_match = re.search(r'<function=([^>]+)>(.*?)</function>', text, re.DOTALL)
    if xml_match:
        action_name = xml_match.group(1).strip()
        params_str = xml_match.group(2).strip()
        params = {}
        if params_str:
            try:
                params = json.loads(params_str)
            except:
                pass
        return {"action": action_name, "params": params}
        
    return None

def execute_action(action_data, api_key: str = None):
    """
    Maps the AI's JSON output to piko.py functions.
    """
    action = action_data.get("action")
    params = action_data.get("params", {})
    
    # Inject provided API Key into all calls, fallback to module-level piko_key (from .env)
    target_key = api_key or piko_key
    if not target_key:
        return {"error": "Missing Piko API Key. Please provide it or set PIKO_API_KEY in .env"}
        
    params["api_key"] = target_key
    
    mapping = {
        "verify_login": piko.verify_login,
        "get_connectors": piko.get_connectors,
        "get_databases": piko.get_databases,
        "create_form": piko.create_form,
        "update_form": piko.update_form,
        "list_forms": piko.list_forms,
        "get_form_snippets": piko.get_form_snippets,
        "create_auth_preset": piko.create_auth_preset,
        "get_auth_preset": piko.get_auth_preset,
        "get_auth_snippets": piko.get_auth_snippets,
        "list_auth_presets": piko.list_auth_presets,
        "create_system": piko.create_system,
        "list_systems": piko.list_systems
    }
    
    if action in mapping:
        func = mapping[action]
        print(f"\n[PIKO EXECUTION] Running {action}...")
        result = func(**params)
        return result
    else:
        return {"error": f"Action '{action}' is not supported."}

def piko_chat():
    messages = [{"role": "system", "content": get_system_prompt()}]
    
    print("--- PIKO AI AGENT ACTIVE ---")
    print("Ask Piko to create forms, setup auth, or manage connectors.")
    
    while True:
        user_input = input("\nYou: ")
        if user_input.lower() in ["exit", "quit"]:
            break
        
        # Build contextual prompt with only relevant chunks
        contextual_prompt = build_contextual_prompt(user_input)
        messages[0] = {"role": "system", "content": contextual_prompt}
        
        messages.append({"role": "user", "content": user_input})
            
        try:
            # The agent can auto-loop until it needs user input
            while True:
                # TRIMMING: Keep system prompt + last 10 messages to maintain speed
                current_messages = [messages[0]] + messages[-20:] if len(messages) > 21 else messages

                completion = client.chat.completions.create(
                    model=os.getenv("PIKO_MODEL"),
                    messages=current_messages,
                    temperature=MAIN_MODEL_TEMP,
                    top_p=MAIN_MODEL_TOP_P,
                    max_tokens=MAIN_MODEL_MAX_TOKENS,
                    stream=False
                )
                
                response_text = completion.choices[0].message.content
                # Handle reasoning content if supported/present
                reasoning_text = getattr(completion.choices[0].message, "reasoning_content", "")
                if reasoning_text:
                    print(f"\n[REASONING]: {reasoning_text}\n")
                
                print(response_text, end="", flush=True)
                
                print() # New line after completion
                messages.append({"role": "assistant", "content": response_text})

                try:
                    action_data = extract_action(response_text)
                    if not action_data:
                        raise ValueError("No action found")
                except (json.JSONDecodeError, ValueError):
                    # If it's not JSON, assume it's just trying to chat
                    print("\n[PIKO]: " + response_text)
                    observation = "Error: Expected JSON format. Please output valid JSON actions."
                    messages.append({"role": "user", "content": observation})
                    continue
                
                action = action_data.get("action")
                
                # If Piko explicitly asks for clarification, break the auto-loop and ask the user
                if action == "ask_clarification":
                    print(f"\n[PIKO]: {action_data['params'].get('question')}")
                    break
                    
                # If Piko signifies the task is complete, break the auto-loop
                if action == "task_complete":
                    print(f"\n[PIKO]: {action_data.get('params', {}).get('message', 'Task completed successfully.')}")
                    break
                
                # 1. Execute the mapped backend action
                result = execute_action(action_data)
                
                # 2. Display the result cleanly using our formatting helpers if applicable
                if result is not None:
                    if action in ["get_form_snippets", "get_auth_snippets"]:
                        piko.display_snippets(result)
                    elif action == "list_forms":
                        piko.display_forms(result)
                    elif action == "get_databases":
                        piko.display_databases(result)
                    else:
                        print(f"\n[RESULT]: {json.dumps(result, indent=2) if isinstance(result, (dict, list)) else result}")
                
                # 3. Feed the result back to model so it can evaluate and decide the next step
                observation = f"Action {action} completed. Result: {result}"
                messages.append({"role": "user", "content": observation})
                
        except Exception as e:
            if "incomplete chunked read" in str(e).lower():
                print("\n[NETWORK ERROR]: Connection was interrupted by the server. Retrying...")
                # We don't break here, the outer loop will allow the user to try again
            else:
                print(f"An error occurred: {e}")

if __name__ == "__main__":
    piko_chat()