static

model/conversation_chain_singleton.py

@@ -1,35 +1,77 @@
+"""
+Module: conversation_chain_singleton
+
+This module provides a singleton class, ConversationChainSingleton, for managing a conversation chain instance.
+
+Dependencies:
+- langchain.memory: Module providing memory functionalities for conversation chains.
+- langchain.chains: Module providing conversation chain functionalities.
+- langchain.llms: Module providing language model functionalities, particularly from HuggingFaceHub.
+
+Classes:
+- ConversationChainSingleton: A singleton class for managing a conversation chain instance.
+"""
+
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationChain
 from langchain.llms import HuggingFaceHub
 
 class ConversationChainSingleton:
+    """
+    A singleton class for managing a conversation chain instance.
+
+    Attributes:
+    - _instance: Private attribute holding the singleton instance.
+    - conversation_chain: The conversation chain instance.
+
+    Methods:
+    - __new__(cls, *args, **kwargs): Creates a new instance of the ConversationChainSingleton class.
+    - get_conversation_chain(self): Returns the conversation chain instance.
+
+    Static Methods:
+    - get_conversation_chain(): Creates and returns a conversational retrieval chain and a language model.
+    """
+
     _instance = None
 
     def __new__(cls, *args, **kwargs):
+        """
+        Create a new instance of the ConversationChainSingleton class if it doesn't exist.
+
+        Returns:
+        - ConversationChainSingleton: The singleton instance.
+        """
         if not cls._instance:
             cls._instance = super(ConversationChainSingleton, cls).__new__(cls)
             # Initialize your conversation chain here
-            cls._instance.conversation_chain = get_conversation_chain()
+            cls._instance.conversation_chain = cls.get_conversation_chain()
         return cls._instance
 
     def get_conversation_chain(self):
-        return self.conversation_chain
-
-
-    def get_conversation_chain( ):
         """
-        Create a conversational retrieval chain and a language model.
+        Get the conversation chain instance.
 
+        Returns:
+        - ConversationChain: The conversation chain instance.
         """
+        return self.conversation_chain
+
+@staticmethod
+def get_conversation_chain():
+    """
+    Create a conversational retrieval chain and a language model.
+
+    Returns:
+    - ConversationChain: The conversation chain instance.
+    """
+    llm = HuggingFaceHub(
+        repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
+        model_kwargs={"max_length": 1048, "temperature": 0.2, "max_new_tokens": 256, "top_p": 0.95, "repetition_penalty": 1.0},
+    )
+    # llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613")
 
-        llm = HuggingFaceHub(
-            repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
-            model_kwargs={"max_length": 1048, "temperature":0.2, "max_new_tokens":256, "top_p":0.95, "repetition_penalty":1.0},
-        )
-        # llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613")
-
-        memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-        conversation_chain = ConversationChain(
-            llm=llm, verbose=True, memory=memory
-        )
-        return conversation_chain
+    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+    conversation_chain = ConversationChain(
+        llm=llm, verbose=True, memory=memory
+    )
+    return conversation_chain

model/custom_agent.py

@@ -1,17 +1,41 @@
-# custom_agent.py
-import os
-import base64
-import io
-import requests
+"""
+Module: custom_agent
+
+This module provides a custom class, CustomHfAgent, for interacting with the Hugging Face model API.
+
+Dependencies:
+- time: Standard Python time module for time-related operations.
+- requests: HTTP library for making requests.
+- transformers: Hugging Face's transformers library for NLP tasks.
+- utils.logger: Custom logger module for logging responses.
+
+Classes:
+- CustomHfAgent: A custom class for interacting with the Hugging Face model API.
+"""
+
 import time
+import requests
 from transformers import Agent
 from utils.logger import log_response
 
-import time  
-import torch
-
 class CustomHfAgent(Agent):
+    """A custom class for interacting with the Hugging Face model API."""
+
     def __init__(self, url_endpoint, token, chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None):
+        """
+        Initialize the CustomHfAgent.
+
+        Args:
+        - url_endpoint (str): The URL endpoint for the Hugging Face model API.
+        - token (str): The authentication token required to access the API.
+        - chat_prompt_template (str): Template for chat prompts.
+        - run_prompt_template (str): Template for run prompts.
+        - additional_tools (list): Additional tools for the agent.
+        - input_params (dict): Additional parameters for input.
+
+        Returns:
+        - None
+        """
         super().__init__(
             chat_prompt_template=chat_prompt_template,
             run_prompt_template=run_prompt_template,
@@ -22,6 +46,16 @@ class CustomHfAgent(Agent):
         self.input_params = input_params
 
     def generate_one(self, prompt, stop):
+        """
+        Generate one response from the Hugging Face model.
+
+        Args:
+        - prompt (str): The prompt to generate a response for.
+        - stop (list): A list of strings indicating where to stop generating text.
+
+        Returns:
+        - str: The generated response.
+        """
         headers = {"Authorization": self.token}
         max_new_tokens = self.input_params.get("max_new_tokens", 192)
         parameters = {"max_new_tokens": max_new_tokens, "return_full_text": False, "stop": stop, "padding": True, "truncation": True}
@@ -29,12 +63,17 @@ class CustomHfAgent(Agent):
             "inputs": prompt,
             "parameters": parameters,
         }
-        response = requests.post(self.url_endpoint, json=inputs, headers=headers)
-
+        print(inputs)
+        try:
+            response = requests.post(self.url_endpoint, json=inputs, headers=headers, timeout=300)
+        except requests.Timeout:
+            pass
+        except requests.ConnectionError:
+            pass
         if response.status_code == 429:
             log_response("Getting rate-limited, waiting a tiny bit before trying again.")
             time.sleep(1)
-            return self._generate_one(prompt)
+            return self.generate_one(prompt, stop)
         elif response.status_code != 200:
             raise ValueError(f"Errors {inputs} {response.status_code}: {response.json()}")
         log_response(response)