chatbot.py

import logging
from typing import List, Dict, Any, Optional
import ollama
import chromadb
from chromadb.utils import embedding_functions
from pydantic_settings import BaseSettings
import asyncio
import os
from enum import Enum
import openai
from openai import AsyncOpenAI
import anthropic
import google.generativeai as genai
import time
import subprocess
import requests
import re
import markdown

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Provider(str, Enum):
    """
    Supported AI model providers.
    Each provider requires specific configuration and API keys.
    """
    OLLAMA = "ollama"
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    GOOGLE = "google"
    HUGGINGFACE = "huggingface"

class ChatbotConfig(BaseSettings):
    """
    Configuration settings for the chatbot.
    Includes model provider selection, API keys, and behavior parameters.
    
    Attributes:
        PROVIDER: The AI model provider to use
        MODEL_NAME: Name of the specific model to use
        API_KEY: API key for cloud providers
        KNOWLEDGE_BASE_DIR: Directory for storing knowledge base
        COLLECTION_NAME: Name of the ChromaDB collection
        MAX_HISTORY_LENGTH: Maximum number of messages to keep in history
        CONTEXT_LENGTH: Number of previous messages to include as context
        LOG_LEVEL: Logging level
        SYSTEM_PROMPT_TEMPLATE: System prompt template
    """
    PROVIDER: str = Provider.OLLAMA
    MODEL_NAME: str = "gemma:2b"
    API_KEY: str = ""
    KNOWLEDGE_BASE_DIR: str = "./knowledge_base"
    COLLECTION_NAME: str = "company_knowledge"
    MAX_HISTORY_LENGTH: int = 10
    CONTEXT_LENGTH: int = 3
    LOG_LEVEL: str = "INFO"

    # Default System Prompt Template (can be overridden)
    # System prompt template updated for smart extraction, flexible Q&A, and plain text output (no markdown)
    DEFAULT_SYSTEM_PROMPT_TEMPLATE: str = """You are a professional and friendly company assistant dedicated to delivering accurate customer service based solely on the provided company knowledge base.

Instructions:
- Use only the information in {context_text} for your answers.
- Extract and present relevant facts (such as product names, categories, prices, specifications, features, etc.) from the provided information, even if it is in table or CSV format.
- Answer flexibly: if the user asks about price, features, comparisons, or availability, extract and present the answer clearly using the data you have.
- When listing products, include only those strictly relevant to the user's request (e.g., only smartphones for a smartphone query).
- If you cannot find enough information to answer, reply: "I don't have this information in my knowledge base."
- Keep responses concise and clear. Do NOT use any markdown formatting (no bullet points, no code, no tables, etc.). Answer with plain text lists or sentences only.

Examples:
Q: "What smartphones are available?"
A: Google Pixel 8, Samsung Galaxy S24 Ultra, Apple iPhone 15 Pro.

Q: "How much is the Apple iPhone 15 Pro?"
A: The price of Apple iPhone 15 Pro is 1199.99.

Q: "Compare iPhone 15 Pro and Samsung Galaxy S24 Ultra."
A: Apple iPhone 15 Pro: 6.7-inch display, A17 Pro chip, 48MP camera, 1199.99. Samsung Galaxy S24 Ultra: 6.8-inch display, Snapdragon 8 Gen 3, 200MP camera, 1299.99.
"""

    SYSTEM_PROMPT_TEMPLATE: str = DEFAULT_SYSTEM_PROMPT_TEMPLATE

    class Config:
        env_prefix = "CHATBOT_"

    def __post_init__(self):
        self.load_from_env()
        self.LOG_LEVEL = os.environ.get("LOG_LEVEL", self.LOG_LEVEL)
        self.SYSTEM_PROMPT_TEMPLATE = os.environ.get("SYSTEM_PROMPT_TEMPLATE", self.SYSTEM_PROMPT_TEMPLATE)

    def update_provider(self, provider: str, model_name: str, api_key: Optional[str] = None):
        """
        Update the model provider configuration.
        
        Args:
            provider: Name of the provider to use
            model_name: Name of the model to use
            api_key: API key for cloud providers (if required)
        """
        self.PROVIDER = provider
        self.MODEL_NAME = model_name
        self.API_KEY = api_key

class KnowledgeBase:
    """Knowledge base management using ChromaDB"""

    def __init__(self, persist_directory: str, collection_name: str):
        self.client = chromadb.Client(chromadb.Settings(
            persist_directory=persist_directory,
            anonymized_telemetry=False
        ))
        
        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction()
        
        # Recreate collection each time to ensure fresh start
        try:
            self.client.delete_collection(collection_name)
        except Exception:
            pass
            
        self.collection = self.client.create_collection(
            name=collection_name,
            embedding_function=self.embedding_function
        )
        
        logger.info("Knowledge base initialized successfully")

    def add_documents(self, documents: List[Dict[str, Any]]) -> bool:
        """Add documents to knowledge base"""
        try:
            texts = [doc["text"] for doc in documents]
            metadatas = [doc.get("metadata", {}) for doc in documents]
            # Use unique ID from metadata if available, otherwise fallback to source
            ids = [doc["metadata"].get("id", doc["metadata"]["source"]) for doc in documents]
            
            self.collection.add(
                documents=texts,
                metadatas=metadatas,
                ids=ids
            )
            logger.info(f"Added {len(documents)} documents to knowledge base")
            return True
            
        except Exception as e:
            logger.error(f"Error adding documents: {str(e)}")
            raise

    def search(self, query: str, n_results: int = 3) -> List[Dict[str, Any]]:
        """Search for relevant documents with metadata"""
        try:
            results = self.collection.query(
                query_texts=[query],
                n_results=n_results
            )
            
            documents = []
            for i in range(len(results['documents'][0])):
                documents.append({
                    'text': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'source': results['ids'][0][i]
                })
                
            logger.info(f"Found {len(documents)} documents for query: '{query}'")
            return documents
            
        except Exception as e:
            logger.error(f"Error searching documents: {str(e)}")
            raise

    def reload_documents(self, documents: List[Dict[str, Any]]) -> bool:
        """Clear existing collection and add new documents."""
        try:
            # Clear existing documents (efficiently)
            # Get all IDs first
            existing_ids = self.collection.get()['ids']
            if existing_ids:
                 self.collection.delete(ids=existing_ids)
                 logger.info(f"Cleared {len(existing_ids)} documents from collection.")
            
            # Add new documents
            return self.add_documents(documents)
        except Exception as e:
            logger.error(f"Error reloading documents: {str(e)}")
            return False

def load_knowledge_files(directory: str = "knowledge") -> List[Dict[str, Any]]:
    """Load knowledge from txt files"""
    documents = []
    
    if not os.path.exists(directory):
        os.makedirs(directory)
        logger.info(f"Created directory {directory}")
        return documents
    
    for filename in os.listdir(directory):
        if filename.endswith(".txt"):
            filepath = os.path.join(directory, filename)
            try:
                with open(filepath, 'r', encoding='utf-8') as f:
                    text = f.read().strip()
                    documents.append({
                        "text": text,
                        "metadata": {
                            "source": filename,
                            "category": filename.replace(".txt", "")
                        }
                    })
                logger.info(f"Loaded file: {filename}")
            except Exception as e:
                logger.error(f"Error reading file {filename}: {str(e)}")
    
    return documents

class ModelManager:
    def __init__(self, model_name: str = "gemma:2b"):
        self.model_name = model_name
        self.logger = logging.getLogger(__name__)
        self.max_retries = 3
        self.retry_delay = 2  # seconds
        self.ollama_url = "http://localhost:11434"
        
    def check_ollama_service(self) -> bool:
        """Check if Ollama service is running."""
        try:
            response = requests.get(f"{self.ollama_url}/api/tags")
            return response.status_code == 200
        except requests.exceptions.ConnectionError:
            return False
    
    def start_ollama_service(self) -> bool:
        """Start Ollama service if it's not running."""
        try:
            subprocess.Popen(["ollama", "serve"], 
                           stdout=subprocess.PIPE,
                           stderr=subprocess.PIPE)
            # Wait for service to start
            for _ in range(5):
                if self.check_ollama_service():
                    return True
                time.sleep(1)
            return False
        except FileNotFoundError:
            self.logger.error("Ollama binary not found. Please install Ollama first.")
            return False
        except Exception as e:
            self.logger.error(f"Failed to start Ollama service: {str(e)}")
            return False
    
    def check_model_availability(self) -> bool:
        """Check if the required model is available locally."""
        try:
            response = requests.get(f"{self.ollama_url}/api/tags")
            if response.status_code == 200:
                models = response.json().get("models", [])
                return any(model["name"] == self.model_name for model in models)
            return False
        except Exception as e:
            self.logger.error(f"Failed to check model availability: {str(e)}")
            return False
    
    def pull_model(self) -> bool:
        """Pull the required model if not available locally."""
        try:
            response = requests.post(
                f"{self.ollama_url}/api/pull",
                json={"name": self.model_name}
            )
            return response.status_code == 200
        except Exception as e:
            self.logger.error(f"Failed to pull model: {str(e)}")
            return False
    
    def initialize(self) -> Optional[str]:
        """Initialize the model with proper error handling and retries."""
        error_message = None
        
        # Step 1: Check and start Ollama service
        if not self.check_ollama_service():
            self.logger.info("Ollama service not running. Attempting to start...")
            if not self.start_ollama_service():
                error_message = "Failed to start Ollama service. Please ensure Ollama is installed correctly."
                self.logger.error(error_message)
                return error_message
        
        # Step 2: Check model availability
        if not self.check_model_availability():
            self.logger.info(f"Model {self.model_name} not found locally. Attempting to pull...")
            if not self.pull_model():
                error_message = f"Failed to pull model {self.model_name}. Please check your internet connection and try again."
                self.logger.error(error_message)
                return error_message
        
        # Step 3: Verify model is working
        for attempt in range(self.max_retries):
            try:
                response = requests.post(
                    f"{self.ollama_url}/api/generate",
                    json={
                        "model": self.model_name,
                        "prompt": "test",
                        "stream": False
                    }
                )
                if response.status_code == 200:
                    self.logger.info(f"Model {self.model_name} initialized successfully")
                    return None
                
                error_message = f"Model test failed (Attempt {attempt + 1}/{self.max_retries})"
                self.logger.warning(error_message)
                time.sleep(self.retry_delay)
                
            except Exception as e:
                error_message = f"Error testing model: {str(e)}"
                self.logger.error(error_message)
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay)
        
        return error_message
    
    def generate(self, prompt: str, **kwargs) -> Optional[str]:
        """Generate response from the model with error handling."""
        for attempt in range(self.max_retries):
            try:
                response = requests.post(
                    f"{self.ollama_url}/api/generate",
                    json={
                        "model": self.model_name,
                        "prompt": prompt,
                        "stream": False,
                        **kwargs
                    }
                )
                
                if response.status_code == 200:
                    return response.json().get("response")
                
                self.logger.warning(f"Generation failed (Attempt {attempt + 1}/{self.max_retries})")
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay)
                
            except Exception as e:
                self.logger.error(f"Error during generation: {str(e)}")
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay)
        
        return None

class ResponseFormatter:
    """Format and enhance model responses."""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.code_pattern = re.compile(r'```(\w+)?\n(.*?)\n```', re.DOTALL)
        self.list_pattern = re.compile(r'^\s*[-*]\s+(.+)$', re.MULTILINE)
        self.number_list_pattern = re.compile(r'^\s*\d+\.\s+(.+)$', re.MULTILINE)
    
    def format_code_blocks(self, text: str) -> str:
        """Format code blocks with syntax highlighting."""
        def replace_code(match):
            language = match.group(1) or ''
            code = match.group(2)
            return f'<pre><code class="language-{language}">{code}</code></pre>'
        
        return self.code_pattern.sub(replace_code, text)
    
    def format_lists(self, text: str) -> str:
        """Format bullet points and numbered lists."""
        # Format bullet points
        text = self.list_pattern.sub(r'<li>\1</li>', text)
        text = re.sub(r'(<li>.*?</li>\s*)+', r'<ul>\g<0></ul>', text, flags=re.DOTALL)
        
        # Format numbered lists
        text = self.number_list_pattern.sub(r'<li>\1</li>', text)
        text = re.sub(r'(<li>.*?</li>\s*)+', r'<ol>\g<0></ol>', text, flags=re.DOTALL)
        
        return text
    
    def format_tables(self, text: str) -> str:
        """Convert markdown tables to HTML tables."""
        lines = text.split('\n')
        in_table = False
        table_lines = []
        formatted_lines = []
        
        for line in lines:
            if line.strip().startswith('|') and line.strip().endswith('|'):
                if not in_table:
                    in_table = True
                table_lines.append(line)
            else:
                if in_table:
                    # Convert collected table lines to HTML
                    if len(table_lines) > 2:  # Valid table needs header and separator
                        formatted_lines.append(self.convert_table_to_html(table_lines))
                    table_lines = []
                    in_table = False
                formatted_lines.append(line)
        
        # Handle table at end of text
        if in_table and len(table_lines) > 2:
            formatted_lines.append(self.convert_table_to_html(table_lines))
        
        return '\n'.join(formatted_lines)
    
    def convert_table_to_html(self, table_lines: List[str]) -> str:
        """Convert markdown table lines to HTML table."""
        # Remove outer pipes and split cells
        rows = [
            [cell.strip() for cell in line.strip('|').split('|')]
            for line in table_lines if not all(c in '|-' for c in line)
        ]
        
        if not rows:
            return ''
        
        html = ['<table class="markdown-table">']
        
        # Add header
        html.append('<thead><tr>')
        for cell in rows[0]:
            html.append(f'<th>{cell}</th>')
        html.append('</tr></thead>')
        
        # Add body
        if len(rows) > 1:
            html.append('<tbody>')
            for row in rows[1:]:
                html.append('<tr>')
                for cell in row:
                    html.append(f'<td>{cell}</td>')
                html.append('</tr>')
            html.append('</tbody>')
        
        html.append('</table>')
        return '\n'.join(html)
    
    def format_response(self, text: str) -> str:
        """Format the complete response with all enhancements."""
        try:
            # Step 1: Format code blocks (before markdown to prevent interference)
            text = self.format_code_blocks(text)
            
            # Step 2: Format tables
            text = self.format_tables(text)
            
            # Step 3: Convert markdown to HTML
            text = markdown.markdown(text)
            
            # Step 4: Format lists (after markdown for better list detection)
            text = self.format_lists(text)
            
            return text
            
        except Exception as e:
            self.logger.error(f"Error formatting response: {str(e)}")
            return text  # Return original text if formatting fails

class Chatbot:
    """Main chatbot class"""

    def __init__(self, config: Optional[ChatbotConfig] = None):
        self.config = config or ChatbotConfig()
        self.conversation_history: List[Dict[str, str]] = []
        self.knowledge_base = KnowledgeBase(
            persist_directory=self.config.KNOWLEDGE_BASE_DIR,
            collection_name=self.config.COLLECTION_NAME
        )
        self.model_client = None
        self.response_formatter = ResponseFormatter()
        self.setup_client()

    def setup_client(self):
        """Setup appropriate client based on provider"""
        if self.config.PROVIDER == Provider.OLLAMA:
            self.model_client = ollama.Client()
        elif self.config.PROVIDER == Provider.OPENAI:
            self.model_client = AsyncOpenAI(api_key=self.config.API_KEY)
        elif self.config.PROVIDER == Provider.ANTHROPIC:
            self.model_client = anthropic.AsyncAnthropic(api_key=self.config.API_KEY)
        elif self.config.PROVIDER == Provider.GOOGLE:
            genai.configure(api_key=self.config.API_KEY)
            self.model_client = genai
        elif self.config.PROVIDER == Provider.HUGGINGFACE:
            # Hugging Face Inference API: just save config, client not needed
            self.model_client = {
                "api_key": self.config.API_KEY,
                "model": self.config.MODEL_NAME,
                "endpoint": f"https://api-inference.huggingface.co/models/{self.config.MODEL_NAME}"
            }

    def get_system_prompt_template(self) -> str:
        """Returns the current system prompt template."""
        return self.config.SYSTEM_PROMPT_TEMPLATE

    def set_system_prompt_template(self, template: str):
        """Sets a new system prompt template."""
        self.config.SYSTEM_PROMPT_TEMPLATE = template

    async def get_response(self, messages: List[Dict[str, str]]) -> str:
        """Get response from the appropriate provider"""
        try:
            if self.config.PROVIDER == Provider.OLLAMA:
                response = self.model_client.chat(
                    model=self.config.MODEL_NAME,
                    messages=messages
                )
                return response['message']['content']
            elif self.config.PROVIDER == Provider.OPENAI:
                response = await self.model_client.chat.completions.create(
                    model=self.config.MODEL_NAME,
                    messages=messages
                )
                return response.choices[0].message.content
            elif self.config.PROVIDER == Provider.ANTHROPIC:
                messages_text = "\n\n".join([f"{m['role']}: {m['content']}" for m in messages])
                response = await self.model_client.messages.create(
                    model=self.config.MODEL_NAME,
                    max_tokens=1024,
                    messages=[{"role": "user", "content": messages_text}]
                )
                return response.content[0].text
            elif self.config.PROVIDER == Provider.GOOGLE:
                model = self.model_client.GenerativeModel(self.config.MODEL_NAME)
                response = model.generate_content([m['content'] for m in messages])
                return response.text
            elif self.config.PROVIDER == Provider.HUGGINGFACE:
                import httpx
                api_key = self.model_client["api_key"]
                endpoint = self.model_client["endpoint"]
                headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
                # Hugging Face Inference API expects {'inputs': ...}
                # We'll concatenate messages for chat-like models
                prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
                payload = {"inputs": prompt}
                async with httpx.AsyncClient(timeout=60.0) as client:
                    resp = await client.post(endpoint, headers=headers, json=payload)
                    resp.raise_for_status()
                    data = resp.json()
                    # For text-generation models, output is usually {'generated_text': ...} or a list of such dicts
                    if isinstance(data, list) and "generated_text" in data[0]:
                        return data[0]["generated_text"]
                    elif isinstance(data, dict) and "generated_text" in data:
                        return data["generated_text"]
                    elif isinstance(data, dict) and "error" in data:
                        raise RuntimeError(f"Hugging Face API error: {data['error']}")
                    else:
                        # Try to get the first string output
                        return str(data)
        except Exception as e:
            logger.error(f"Error getting response from {self.config.PROVIDER}: {str(e)}")
            raise

    def add_knowledge(self, documents: List[Dict[str, Any]]) -> bool:
        """Add documents to knowledge base"""
        self.knowledge_base.add_documents(documents)
        return True

    async def reload_knowledge(self) -> bool:
        """Reload knowledge from files into the vector store."""
        logger.info("Reloading knowledge base...")
        try:
            documents = load_knowledge_files()
            if not documents:
                 logger.warning("No knowledge files found to reload.")
                 # Decide if we should clear the collection even if no files are found
                 # For now, let's clear it to reflect the state of the knowledge folder
                 existing_ids = self.knowledge_base.collection.get()['ids']
                 if existing_ids:
                      self.knowledge_base.collection.delete(ids=existing_ids)
                      logger.info(f"Cleared {len(existing_ids)} documents as no new files found.")
                 return True
            
            # Use the reload_documents method for efficiency
            reloaded = self.knowledge_base.reload_documents(documents)
            if reloaded:
                 logger.info("Knowledge base reloaded successfully.")
                 return True
            else:
                 logger.error("Failed to reload documents into knowledge base.")
                 return False
        except Exception as e:
            logger.error(f"Error during knowledge reload: {str(e)}")
            return False

    async def process_message(self, message: str) -> Dict[str, Any]:
        """Process incoming user message"""
        try:
            logger.info(f"Received user message: '{message}'")
            
            # Search for relevant context
            context_docs = self.knowledge_base.search(
                message, 
                n_results=self.config.CONTEXT_LENGTH
            )
            
            # Context is passed to LLM as-is; LLM is responsible for extracting relevant info (no keyword filtering)
            # Format context with source information
            context_text = "\n\n".join([
                f"From {doc['source']}:\n{doc['text']}" 
                for doc in context_docs
            ])

            # --- LOGGING ADDED ---
            logger.info(f"Retrieved Context for '{message}':\n{context_text if context_text else '<<-- No context found -->>'}")
            # --- END LOGGING ---

            # Prepare system prompt using the template from config
            try:
                system_prompt = self.config.SYSTEM_PROMPT_TEMPLATE.format(context_text=context_text)
            except KeyError as e:
                logger.error(f"Error formatting system prompt template. Missing key: {e}. Using default context placeholder.")
                # Fallback if formatting fails (e.g., template changed and key is missing)
                system_prompt = self.config.SYSTEM_PROMPT_TEMPLATE.replace("{context_text}", context_text)

            # Prepare messages
            messages = [
                {"role": "system", "content": system_prompt},
                *self.conversation_history,
                {"role": "user", "content": message}
            ]
            
            # Get model response
            assistant_response = await self.get_response(messages)
            
            # Update conversation history
            self.conversation_history.extend([
                {"role": "user", "content": message},
                {"role": "assistant", "content": assistant_response}
            ])
            
            # Limit history length
            if len(self.conversation_history) > self.config.MAX_HISTORY_LENGTH:
                self.conversation_history = self.conversation_history[-self.config.MAX_HISTORY_LENGTH:]
            
            # Format the response
            formatted_response = self.response_formatter.format_response(assistant_response)
            
            return {
                "status": "success",
                "response": formatted_response,
                "raw_response": assistant_response,
                "sources": [doc['source'] for doc in context_docs]
            }
            
        except Exception as e:
            logger.error(f"Error processing message: {str(e)}")
            return {
                "status": "error",
                "error": str(e)
            }

    def clear_history(self):
        """Clear conversation history"""
        self.conversation_history = []
        logger.info("Conversation history cleared")

    def initialize(self) -> Optional[str]:
        """Initialize the chatbot. (Model client is setup in __init__)"""
        # Original call caused error as client objects don't have .initialize()
        # return self.model_client.initialize() 
        
        # Placeholder: Add any necessary startup checks here if needed,
        # e.g., specifically for Ollama model verification/pulling.
        # For now, assume success if __init__ completed.
        logger.info("Chatbot initialize method called.")
        return None

async def main():
    # Initialize bot
    bot = Chatbot()
    
    # Load knowledge files
    documents = load_knowledge_files()
    if documents:
        bot.add_knowledge(documents)
    else:
        logger.warning("No knowledge files found in knowledge directory")
    
    print("\nChatbot is ready. Type 'exit' to quit.\n")
    
    while True:
        try:
            user_input = input("You: ")
            
            if user_input.lower() in ['выход', 'exit', 'quit']:
                print("Goodbye!")
                break
            
            response = await bot.process_message(user_input)
            
            if response["status"] == "success":
                print(f"\nBot: {response['response']}")
                print(f"Sources used: {', '.join(response['sources'])}\n")
            else:
                print(f"\nError: {response['error']}\n")
                
        except KeyboardInterrupt:
            print("\nOperation terminated")
            break
        except Exception as e:
            print(f"\nAn error occurred: {str(e)}\n")

if __name__ == "__main__":
    asyncio.run(main())