Upload 4 files

src/chatbot.py

@@ -0,0 +1,72 @@
+import gradio as gr
+from langchain.chains import ConversationChain
+from .memory import EnhancedInMemoryHistory
+from .models import route_llm, prompt
+
+# Function to process input and generate a response
+def process_input(user_input, session_id='1'):
+    """
+    Processes the user input and generates a response using the conversation chain.
+    
+    Parameters:
+        user_input (str): The user's input message.
+        session_id (str): The session ID for the chat (default is "1").
+    
+    Returns:
+        str: The generated response from the chatbot.
+    """
+    memory = get_by_session_id(session_id)
+    
+    if user_input.lower() == 'exit':
+        return "Exiting the chat session."
+
+    llm = route_llm(user_input)
+
+    conversation_chain = ConversationChain(
+        llm=llm,
+        prompt=prompt,
+        memory=memory,
+        input_key='input',
+        verbose=True
+    )
+
+    response = conversation_chain.run({"input": user_input})
+    memory.save_context({'input': user_input}, response)
+    
+    return response
+
+# Function for Gradio interface
+def chatbot_interface(user_input, chat_history=None, session_id="1"):
+    """
+    Interface function for Gradio to handle input and output between the user and the chatbot.
+    
+    Parameters:
+        user_input (str): The user's input message.
+        session_id (str): The session ID for the chat (default is "1").
+        chat_history (list): List of previous chat messages in the format [[user, bot], ...]
+    
+    Returns:
+        list: Updated chat history including the new user and bot messages.
+    """
+    if chat_history is None:
+        chat_history = []
+    
+    if user_input == "":
+        bot_response = "Hi there! How can I help you today?"
+    else:
+        bot_response = process_input(user_input, session_id)
+    
+    chat_history.append([user_input, bot_response])
+    
+    return chat_history
+
+# Gradio launch
+def launch_gradio_interface():
+    gr.Interface(
+        fn=chatbot_interface,
+        inputs=[gr.Textbox(lines=7, label="Your input", placeholder="Type your message here...")],
+        outputs=gr.Chatbot(label="Chat History"),
+        title="AI Chatbot",
+        live=False
+    ).launch()
+

src/memory.py

@@ -0,0 +1,77 @@
+from langchain_core.memory import BaseMemory
+from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
+from langchain_core.pydantic_v1 import Field
+from typing import List
+from datetime import datetime
+
+class EnhancedInMemoryHistory(BaseMemory):
+    """
+    Custom memory class for storing chat history with timestamps.
+
+    Attributes:
+        messages (List[BaseMessage]): A list of messages exchanged between the user and the bot.
+        timestamps (List[datetime]): A list of timestamps when the messages were exchanged.
+    """
+
+    messages: List[BaseMessage] = Field(default_factory=list)
+    timestamps: List[datetime] = Field(default_factory=list)
+
+    @property
+    def memory_variables(self):
+        """Returns a list of memory variables (history) used in the conversation chain."""
+        return ["history"]
+
+    def add_messages(self, messages: List[BaseMessage]):
+        """
+        Adds new messages to the memory and timestamps them.
+
+        Parameters:
+            messages (List[BaseMessage]): A list of messages to add to the memory.
+        """
+        current_time = datetime.now()
+        self.messages.extend(messages)
+        self.timestamps.extend([current_time] * len(messages))
+
+    def get_recent_messages(self, limit: int = 5):
+        """
+        Retrieves the most recent messages.
+
+        Parameters:
+            limit (int): Number of recent messages to retrieve (default is 5).
+
+        Returns:
+            List[BaseMessage]: A list of the most recent messages.
+        """
+        return self.messages[-limit:]
+
+    def clear(self):
+        """Clears all messages and timestamps from the memory."""
+        self.messages = []
+        self.timestamps = []
+
+    def load_memory_variables(self, inputs: dict):
+        """
+        Loads memory variables for the conversation chain.
+
+        Parameters:
+            inputs (dict): Input data for the conversation chain.
+
+        Returns:
+            dict: A dictionary with the conversation history.
+        """
+        return {"history": "\n".join([msg.content for msg in self.messages])}
+
+    def save_context(self, inputs: dict, outputs: str):
+        """
+        Saves the context of the conversation by storing the user input and bot output.
+
+        Parameters:
+            inputs (dict): The user input.
+            outputs (str): The bot's response.
+        """
+        self.add_messages([HumanMessage(content=inputs['input']), AIMessage(content=str(outputs))])
+
+    def clear_memory(self):
+        """Clears the memory (alias for the clear function)."""
+        self.clear()
+

src/models.py

@@ -0,0 +1,42 @@
+from transformers import pipeline
+from langchain.llms import CTransformers
+from langchain.prompts import ChatPromptTemplate
+
+# Text classification model for routing input
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+
+# Load LLM models using CTransformers
+general_llm = CTransformers(
+    model="./llama-2-7b.Q4_K_M.gguf", 
+    model_type="llama",
+    config={'max_new_tokens': 512, 'temperature': 0.7}
+)
+
+medical_llm = CTransformers(
+    model="./BioMistral-7B.Q4_K_M.gguf", 
+    model_type="llama",
+    config={'max_new_tokens': 512, 'temperature': 0.7}
+)
+
+# Prompt template for generating responses
+template = """
+You are a versatile AI assistant that can provide both medical advice and help users with general concerns...
+"""
+
+# Compile the prompt template using LangChain
+prompt = ChatPromptTemplate.from_template(template)
+
+def route_llm(user_input):
+    """
+    Routes user input to the appropriate LLM (medical or general).
+
+    Parameters:
+        user_input (str): The user's input message.
+
+    Returns:
+        CTransformers: The selected LLM model (general or medical).
+    """
+    result = classifier(user_input, ['medical', 'general'])
+    label = result['labels'][0] if 'labels' in result else 'general'
+    return medical_llm if label == 'medical' else general_llm
+