Chris4K/A.I.StateMachine

1 2#chat_service.py3from typing import List, Dict, Any, Optional, Tuple4from datetime import datetime5import logging6from config.config import settings7import asyncio8from io import StringIO  9import pandas as pd10 11logger = logging.getLogger(__name__)12 13class ConversationManager:14    """Manages conversation history and context"""15    def __init__(self):16        self.conversations: Dict[str, List[Dict[str, Any]]] = {}17        self.max_history = 118        19    def add_interaction(20        self,21        session_id: str,22        user_input: str,23        response: str,24        context: Optional[Dict[str, Any]] = None25    ) -> None:26        if session_id not in self.conversations:27            self.conversations[session_id] = []28            29        self.conversations[session_id].append({30            'timestamp': datetime.now().isoformat(),31            'user_input': user_input,32            'response': response,33            'context': context34        })35        36        if len(self.conversations[session_id]) > self.max_history:37            self.conversations[session_id] = self.conversations[session_id][-self.max_history:]38            39    def get_history(self, session_id: str) -> List[Dict[str, Any]]:40        return self.conversations.get(session_id, [])41        42    def clear_history(self, session_id: str) -> None:43        if session_id in self.conversations:44            del self.conversations[session_id]45 46class ChatService:47    def __init__(48        self,49        model_service,50        data_service,51        pdf_service,52        faq_service53    ):54        self.model = model_service.model55        self.tokenizer = model_service.tokenizer56        self.data_service = data_service57        self.pdf_service = pdf_service58        self.faq_service = faq_service59        self.conversation_manager = ConversationManager()60 61    async def search_all_sources(62        self,63        query: str,64        top_k: int = 365    ) -> Dict[str, List[Dict[str, Any]]]:66        """Search across all available data sources"""67        try:68            print("-----------------------------")69            print("starting searches .... ")70            71            # Await the search calls since they're coroutines72            products = await self.data_service.search(query, top_k)73            pdfs = await self.pdf_service.search(query, top_k)74            faqs = await self.faq_service.search_faqs(query, top_k)75 76            results = {77                'products': products or [],78                'documents': pdfs or [],79                'faqs': faqs or []80            }81            82            print("Search results:", results)83            return results84            85        except Exception as e:86            logger.error(f"Error searching sources: {e}")87            return {'products': [], 'documents': [], 'faqs': []}88 89    def construct_system_prompt(self, context: str) -> str:90        """Constructs the system message."""91        return (92            "You are a friendly bot named: Oma Erna, specializing in Bofrost products and content. Use only the context from this prompt. "93            "Return comprehensive German answers. If possible add product IDs from context. Do not make up information. The context is is truth. "94            "Use the following context (product descriptions and information) for answers:\n\n"95            f"{context}\n\n"96        )97 98    def construct_prompt(99            self, 100            user_input: str, 101            context: str, 102            chat_history: List[Tuple[str, str]], 103            max_history_turns: int = 1104        ) -> str:105        """Constructs the full prompt."""106        system_message = self.construct_system_prompt(context)107        prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>"108 109        for user_msg, assistant_msg in chat_history[-max_history_turns:]:110            prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>"111            prompt += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>"112 113        prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|>"114        prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"115 116        return prompt117 118    def build_context(119        self,120        search_results: Dict[str, List[Dict[str, Any]]],121        chat_history: List[Dict[str, Any]]122    ) -> str:123        """Build context for the model from search results and chat history"""124        context_parts = []125        126        # Add relevant products127        if search_results.get('products'):128            products = search_results['products'][:2]  # Limit to top 2 products129            for product in products:130                context_parts.append(131                    f"Produkt: {product['Name']}\n"132                    f"Beschreibung: {product['Description']}\n"133                    f"Preis: {product['Price']}€\n"134                    f"Kategorie: {product['ProductCategory']}"135                )136        137        # Add relevant PDF content138        if search_results.get('documents'):139            docs = search_results['documents'][:2]140            for doc in docs:141                context_parts.append(142                    f"Aus Dokument '{doc['source']}' (Seite {doc['page']}):\n"143                    f"{doc['text']}"144                )145        146        # Add relevant FAQs147        if search_results.get('faqs'):148            faqs = search_results['faqs'][:2]149            for faq in faqs:150                context_parts.append(151                    f"FAQ:\n"152                    f"Frage: {faq['question']}\n"153                    f"Antwort: {faq['answer']}"154                )155        156        # Add recent chat history157        if chat_history:158            print("--- historiy--- ")159            #recent_history = chat_history[-3:]  # Last 3 interactions160            #history_text = "\n".join(161            #    f"User: {h['user_input']}\nAssistant: {h['response']}"162            #    for h in recent_history163            #)164            #context_parts.append(f"Letzte Interaktionen:\n{history_text}")165        166        print("\n\n".join(context_parts))167        return "\n\n".join(context_parts)168 169    async def chat(170        self,171        user_input: str,172        session_id: Any,173        max_length: int = 8000174    ) -> Tuple[str, List[Tuple[str, str]], Dict[str, List[Dict[str, Any]]]]:175        """Main chat method that coordinates the entire conversation flow."""176        try:177            if not isinstance(session_id, str):178                session_id = str(session_id)179    180            chat_history_raw = self.conversation_manager.get_history(session_id)181            chat_history = [182                (entry['user_input'], entry['response']) for entry in chat_history_raw183            ]184    185            search_results = await self.search_all_sources(user_input)186            print(search_results)187            188            context = self.build_context(search_results, chat_history_raw)189            prompt = self.construct_prompt(user_input, context, chat_history)190            response = self.generate_response(prompt, max_length)191    192            self.conversation_manager.add_interaction(193                session_id,194                user_input,195                response,196                {'search_results': search_results}197            )198    199            formatted_history = [200                (entry['user_input'], entry['response']) 201                for entry in self.conversation_manager.get_history(session_id)202            ]203    204            return response, formatted_history, search_results205    206        except Exception as e:207            logger.error(f"Error in chat: {e}")208            raise209 210    def generate_response(211        self,212        prompt: str,213        max_length: int = 1000214    ) -> str:215        """Generate response using the language model"""216        try:217            print(prompt)218            inputs = self.tokenizer(219                prompt,220                return_tensors="pt",221                truncation=True,222                max_length=4096223            ).to(settings.DEVICE)224            225            outputs = self.model.generate(226                **inputs,227                max_length=max_length,228                num_return_sequences=1,229                temperature=0.7,230                top_p=0.9,231                do_sample=True,232                no_repeat_ngram_size=3,233                early_stopping=False234            ) 235 236            input_ids = self.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=4096).to("cpu")237 238            response = self.tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)239 240                # Remove potential repeated assistant text241            response = response.replace("<|assistant|>", "").strip()242            243            return response.strip()244            245        except Exception as e:246            logger.error(f"Error generating response: {e}")247            raise