main.py

import os
from typing import Optional
from openai import AzureOpenAI

# --- Pydantic v2 -----------------------------------
from pydantic import BaseModel, Field, ConfigDict

# =============================================================================
# Configuration
# =============================================================================
client = AzureOpenAI(
    azure_endpoint="",
    api_key="",
    api_version="",
)

# =============================================================================
# Utilities
# =============================================================================
def load_file(path: str, default: str = "") -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    except Exception:
        return default

# =============================================================================
# Pydantic Tool Definitions (extras forbidden + all fields required)
# =============================================================================
class ConsultDocumentExpert(BaseModel):
    """Call the Document Expert for detailed technical questions about the document."""
    model_config = ConfigDict(extra="forbid")  # additionalProperties:false

    # Use Field(...) to force REQUIRED in JSON Schema
    question: str = Field(
        ..., description=(
            "The user's question for Document Expert. "
            "Use 'MODE_SWITCH' for explicit requests to talk to Document Expert (e.g., 'I want to talk to expert', 'Switch to expert mode'). "
            "Otherwise, pass the actual question that needs document expertise."
        )
    )

class ReturnToGeneral(BaseModel):
    """Return control to the General Assistant for non-document questions."""
    model_config = ConfigDict(extra="forbid")  # additionalProperties:false

    question: str = Field(
        ..., description=(
            "The user's question for General Assistant. "
            "Use 'MODE_SWITCH' for explicit requests to talk to General (e.g., 'I want to talk to general', 'Switch back'). "
            "Otherwise, pass the actual question that General should handle."
        )
    )

# =============================================================================
# Chatbot
# =============================================================================
class DocumentChatbot:
    def __init__(self) -> None:
        # Load prompts once
        self.general_prompt = load_file(
            "first_prompt.md",
            default="You are a General Assistant. Answer clearly and route document-heavy questions to the Document Expert tool.",
        )

        self.document_prompt_template = load_file(
            "document_expert_prompt.md",
            default="You are the Document Expert. Use the following document to answer with technical depth:\n\n[The full document content will be inserted here]",
        )

        self.document_content = load_file(
            "document.md",
            default="[Document content not found. Answer based on whatever is available.]",
        )

        self.document_prompt = self.document_prompt_template.replace(
            "[The full document content will be inserted here]",
            self.document_content,
        )

        self.mode = "general"  # or "document"

        # Optional: print schemas once to verify required/additionalProperties
        if os.getenv("DEBUG_SCHEMA") == "1":
            import json

            print("\n[DEBUG] ConsultDocumentExpert schema:")
            print(json.dumps(ConsultDocumentExpert.model_json_schema(), indent=2))

            print("\n[DEBUG] ReturnToGeneral schema:")
            print(json.dumps(ReturnToGeneral.model_json_schema(), indent=2))

    # --------------------------- General LLM ---------------------------------
    def general_llm(self, user_question: str) -> str:
        # Special greeting if we just switched to General mode
        if user_question == "MODE_SWITCH":
            return (
                "Hello! I'm the General Assistant. "
                "I can help with general questions about the system, getting started, "
                "and I'll connect you with the Document Expert for technical details when needed. "
                "How can I help you today?"
            )

        try:
            response = client.chat.completions.create(
                model="GPT-4.1",
                messages=[
                    {"role": "system", "content": self.general_prompt},
                    {"role": "user", "content": user_question},
                ],
                tools=[
                    {
                        "type": "function",
                        "function": {
                            "name": "consult_document_expert",
                            "description": (
                                "Call when: 1) User asks about document technical details, algorithms, formulas. "
                                "2) User explicitly requests to talk to Document Expert (use question='MODE_SWITCH'). "
                                "3) You don't know the answer and it might be in the document (pass the actual question)."
                            ),
                            "parameters": ConsultDocumentExpert.model_json_schema(),
                            "strict": True,
                        },
                    }
                ],
                tool_choice="auto",
                parallel_tool_calls=False,
            )

            message = response.choices[0].message

            # Tool call?
            if getattr(message, "tool_calls", None):
                tool_call = message.tool_calls[0]
                try:
                    # Validate tool args with Pydantic
                    document_request = ConsultDocumentExpert.model_validate_json(
                        tool_call.function.arguments
                    )

                    if document_request.question == "MODE_SWITCH":
                        print(f"\n[ROUTING] → Document Expert: Mode switch requested")
                    else:
                        print(f"\n[ROUTING] → Document Expert: {document_request.question[:50]}...")

                    self.mode = "document"
                    return self.document_expert_llm(document_request.question)

                except Exception as e:
                    print(f"\n[ERROR] Tool call validation failed: {e}")
                    return "Sorry, there was an issue processing your document request."

            return message.content if message.content else ""

        except Exception as e:
            print(f"\n[ERROR] General LLM failed: {e}")
            return "Sorry, I encountered an error processing your question."

    # --------------------------- Document Expert LLM ------------------------------
    def document_expert_llm(self, user_question: str) -> str:
        # Special greeting if we just switched to Document mode
        if user_question == "MODE_SWITCH":
            return (
                "Hello! I'm the Document Expert, specialized in the provided document. "
                "I can help with technical specifications, algorithms, formulas, and implementation details. "
                "What would you like to explore?"
            )

        try:
            response = client.chat.completions.create(
                model="GPT-4.1",
                messages=[
                    {"role": "system", "content": self.document_prompt},
                    {"role": "user", "content": user_question},
                ],
                tools=[
                    {
                        "type": "function",
                        "function": {
                            "name": "return_to_general",
                            "description": (
                                "Call when: 1) User asks non-document questions (pass the actual question). "
                                "2) User explicitly requests to talk to General (use question='MODE_SWITCH')."
                            ),
                            "parameters": ReturnToGeneral.model_json_schema(),
                            "strict": True,
                        },
                    }
                ],
                tool_choice="auto",
                parallel_tool_calls=False,
            )

            message = response.choices[0].message

            # Tool call?
            if getattr(message, "tool_calls", None):
                tool_call = message.tool_calls[0]
                try:
                    return_request = ReturnToGeneral.model_validate_json(
                        tool_call.function.arguments
                    )

                    if return_request.question == "MODE_SWITCH":
                        print(f"\n[ROUTING] → General: Mode switch requested")
                    else:
                        print(f"\n[ROUTING] → General: {return_request.question[:50]}...")

                    self.mode = "general"
                    return self.general_llm(return_request.question)

                except Exception as e:
                    print(f"\n[ERROR] Tool call validation failed: {e}")
                    return "Sorry, there was an issue routing your question."

            return message.content if message.content else ""

        except Exception as e:
            print(f"\n[ERROR] Document Expert failed: {e}")
            return "Sorry, I encountered an error processing your document question."

    # --------------------------- Router --------------------------------------
    def chat(self, user_question: str) -> str:
        if self.mode == "general":
            return self.general_llm(user_question)
        return self.document_expert_llm(user_question)

# =============================================================================
# CLI
# =============================================================================
def main() -> None:
    print("Document Chatbot - Pydantic Structured Outputs (Azure)")
    print("Type 'exit' to quit, 'debug' to show internal state")
    print("-" * 60)

    bot = DocumentChatbot()

    while True:
        mode_label = "General" if bot.mode == "general" else "Document Expert"

        try:
            user_input = input(f"\n[{mode_label}] You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not user_input:
            continue

        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break

        if user_input.lower() == "debug":
            print(f"Current mode: {bot.mode}")
            print(f"General prompt loaded: {'✓' if bot.general_prompt else '✗'}")
            print(f"Document prompt template loaded: {'✓' if bot.document_prompt_template else '✗'}")
            print(f"Document content loaded: {'✓' if bot.document_content else '✗'}")
            continue

        reply = bot.chat(user_input)
        print(f"\nBot: {reply}")

if __name__ == "__main__":
    main()