ai_forward.py

import requests
import json
import re
from typing import Optional, List, Dict, Any

class AIForwarder:
    def __init__(self, api_url: str, system_prompt_file: str = "summary.txt"):
        """
        Initialize the AI Forwarder

        Args:
            api_url: The API endpoint URL
            system_prompt_file: Path to the file containing system prompt
        """
        self.api_url = api_url
        self.system_prompt = self._load_system_prompt(system_prompt_file)

    def _load_system_prompt(self, file_path: str) -> str:
        """Load system prompt from file"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read().strip()
                print(f"✓ System prompt loaded from {file_path}")
                return content
        except FileNotFoundError:
            print(f"⚠ Warning: {file_path} not found. Using default prompt.")
            return self._get_default_prompt()

    def _get_default_prompt(self) -> str:
        """Default system prompt for sheet forwarding"""
        return """You are an AI assistant that helps users access different sheets.
When a user requests a specific sheet number, respond with the format: <forward>NUMBER</forward>
For example:
- User: "I want sheet 3" → Response: "<forward>3</forward>"
- User: "Show me sheet number 5" → Response: "<forward>5</forward>"
- User: "Can I see the second sheet?" → Response: "<forward>2</forward>"

Always extract the sheet number from the user's request and format it properly."""

    def _extract_sheet_number(self, response: str) -> Optional[int]:
        """Extract sheet number from AI response"""
        match = re.search(r'<forward>(\d+)</forward>', response)
        if match:
            return int(match.group(1))
        return None

    def _clean_chat_history(self, chat_history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Remove any existing system messages from chat history

        Args:
            chat_history: Original chat history

        Returns:
            Cleaned chat history without system messages
        """
        cleaned = []
        for message in chat_history:
            # Skip system messages
            if message.get("role") != "system":
                cleaned.append(message)
        return cleaned

    def process_chat_history(
        self,
        chat_history: List[Dict[str, Any]],
        temperature: float = 0.9,
        top_p: float = 0.95,
        max_tokens: Optional[int] = 1000
    ) -> Dict[str, Any]:
        """
        Process chat history and get AI response
        This is the main function to use when importing this module

        Args:
            chat_history: List of chat messages with roles (user/assistant)
                         Example: [
                             {"role": "user", "content": "I want sheet 3"},
                             {"role": "assistant", "content": "Sure!"},
                             {"role": "user", "content": "Show me sheet 5"}
                         ]
            temperature: Sampling temperature
            top_p: Nucleus sampling parameter
            max_tokens: Maximum tokens to generate

        Returns:
            Dictionary with response and extracted sheet number
        """
        # Clean chat history (remove any system messages)
        cleaned_history = self._clean_chat_history(chat_history)

        # Add our system prompt at the beginning
        final_history = [{"role": "system", "content": self.system_prompt}]
        final_history.extend(cleaned_history)

        # Prepare request payload
        payload = {
            "user_input": None,  # No user input, using chat history
            "chat_history": final_history,
            "temperature": temperature,
            "top_p": top_p,
        }

        try:
            # Send request to API
            response = requests.post(
                self.api_url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=30
            )
            response.raise_for_status()

            # Parse response
            result = response.json()
            assistant_response = result.get("assistant_response", "")

            # Extract sheet number if present
            sheet_number = self._extract_sheet_number(assistant_response)

            return {
                "success": True,
                "response": assistant_response,
                "sheet_number": sheet_number,
                "raw_response": result
            }

        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "error": str(e),
                "response": None,
                "sheet_number": None
            }

    def forward_single_message(
        self,
        user_input: str,
        previous_history: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
        """
        Forward a single user message (with optional previous history)

        Args:
            user_input: User's message
            previous_history: Optional previous chat history

        Returns:
            Dictionary with response and sheet number
        """
        chat_history = previous_history.copy() if previous_history else []
        chat_history.append({"role": "user", "content": user_input})

        return self.process_chat_history(chat_history)


# Standalone testing
if __name__ == "__main__":
    # Initialize forwarder
    api_url = "https://dooratre-xx-gpt-5.hf.space/chat"
    forwarder = AIForwarder(api_url)

    print("\n" + "="*50)
    print("AI Sheet Forwarder - Testing Mode")
    print("="*50 + "\n")

    # Test with chat history
    print("Test 1: Processing chat history")
    print("-" * 50)

    test_history = [
        {"role": "user", "content": "Hello!"},
        {"role": "assistant", "content": "Hi! How can I help you?"},
        {"role": "user", "content": "I want sheet number 3"}
    ]

    result = forwarder.process_chat_history(test_history)

    if result["success"]:
        print(f"AI Response: {result['response']}")
        if result["sheet_number"]:
            print(f"✓ Sheet {result['sheet_number']} detected")
    else:
        print(f"✗ Error: {result['error']}")

    print("\n" + "="*50)
    print("Interactive Mode (type 'quit' to exit)")
    print("="*50 + "\n")

    conversation_history = []

    while True:
        try:
            user_input = input("YOU: ").strip()

            if user_input.lower() == 'quit':
                print("Goodbye!")
                break

            if not user_input:
                continue

            # Add user message to history
            conversation_history.append({"role": "user", "content": user_input})

            # Process the conversation
            result = forwarder.process_chat_history(conversation_history)

            if result["success"]:
                print(f"AI: {result['response']}")

                # Add assistant response to history
                conversation_history.append({
                    "role": "assistant",
                    "content": result['response']
                })

                if result["sheet_number"]:
                    print(f"✓ Sheet {result['sheet_number']} requested\n")
            else:
                print(f"✗ Error: {result['error']}\n")

        except KeyboardInterrupt:
            print("\n\nGoodbye!")
            break
        except Exception as e:
            print(f"✗ Unexpected error: {e}\n")