mb-rag 1.1.46__py3-none-any.whl → 1.1.56.post0__py3-none-any.whl

mb_rag/chatbot/chains.py

@@ -1,206 +1,206 @@
-## file for chaining functions in chatbot
-
-from typing import Optional, List, Dict, Any, Union
-from dataclasses import dataclass
-from langchain.schema.output_parser import StrOutputParser
-from mb_rag.chatbot.prompts import invoke_prompt
-from langchain.schema.runnable import RunnableLambda, RunnableSequence
-from mb_rag.utils.extra import check_package
-
-__all__ = ['Chain', 'ChainConfig']
-
-def check_langchain_dependencies() -> None:
-    """
-    Check if required LangChain packages are installed
-    Raises:
-        ImportError: If any required package is missing
-    """
-    if not check_package("langchain"):
-        raise ImportError("LangChain package not found. Please install it using: pip install langchain")
-    if not check_package("langchain_core"):
-        raise ImportError("LangChain Core package not found. Please install it using: pip install langchain-core")
-
-# Check dependencies before importing
-check_langchain_dependencies()
-
-@dataclass
-class ChainConfig:
-    """Configuration for chain operations"""
-    prompt: Optional[str] = None
-    prompt_template: Optional[str] = None
-    input_dict: Optional[Dict[str, Any]] = None
-
-class Chain:
-    """
-    Class to chain functions in chatbot with improved OOP design
-    """
-    def __init__(self, model: Any, config: Optional[ChainConfig] = None, **kwargs):
-        """
-        Initialize chain
-        Args:
-            model: The language model to use
-            config: Chain configuration
-            **kwargs: Additional arguments
-        """
-        self.model = model
-        self._output_parser = StrOutputParser()
-        self._initialize_config(config, **kwargs)
-
-    @classmethod
-    def from_template(cls, model: Any, template: str, input_dict: Dict[str, Any], **kwargs) -> 'Chain':
-        """
-        Create chain from template
-        Args:
-            model: The language model
-            template: Prompt template
-            input_dict: Input dictionary for template
-            **kwargs: Additional arguments
-        Returns:
-            Chain: New chain instance
-        """
-        config = ChainConfig(
-            prompt_template=template,
-            input_dict=input_dict
-        )
-        return cls(model, config, **kwargs)
-
-    @classmethod
-    def from_prompt(cls, model: Any, prompt: str, **kwargs) -> 'Chain':
-        """
-        Create chain from direct prompt
-        Args:
-            model: The language model
-            prompt: Direct prompt
-            **kwargs: Additional arguments
-        Returns:
-            Chain: New chain instance
-        """
-        config = ChainConfig(prompt=prompt)
-        return cls(model, config, **kwargs)
-
-    def _initialize_config(self, config: Optional[ChainConfig], **kwargs) -> None:
-        """Initialize chain configuration"""
-        if config:
-            self.input_dict = config.input_dict
-            if config.prompt_template:
-                self.prompt = invoke_prompt(config.prompt_template, self.input_dict)
-            else:
-                self.prompt = config.prompt
-        else:
-            self.input_dict = kwargs.get('input_dict')
-            if prompt_template := kwargs.get('prompt_template'):
-                self.prompt = invoke_prompt(prompt_template, self.input_dict)
-            else:
-                self.prompt = kwargs.get('prompt')
-
-    @property
-    def output_parser(self) -> StrOutputParser:
-        """Get the output parser"""
-        return self._output_parser
-
-    @staticmethod
-    def _validate_chain_components(prompt: Any, middle_chain: Optional[List] = None) -> None:
-        """
-        Validate chain components
-        Args:
-            prompt: The prompt to validate
-            middle_chain: Optional middle chain to validate
-        Raises:
-            ValueError: If validation fails
-        """
-        if prompt is None:
-            raise ValueError("Prompt is not provided")
-        if middle_chain is not None and not isinstance(middle_chain, list):
-            raise ValueError("middle_chain should be a list")
-
-    def invoke(self) -> Any:
-        """
-        Invoke the chain
-        Returns:
-            Any: Output from the chain
-        Raises:
-            Exception: If prompt is not provided
-        """
-        self._validate_chain_components(self.prompt)
-        chain_output = self.prompt | self.model | self.output_parser
-        return chain_output
-
-    def chain_sequence_invoke(self, 
-                            middle_chain: Optional[List] = None,
-                            final_chain: Optional[RunnableLambda] = None) -> Any:
-        """
-        Chain invoke the sequence
-        Args:
-            middle_chain: List of functions/Prompts/RunnableLambda to chain
-            final_chain: Final chain to run
-        Returns:
-            Any: Output from the chain
-        """
-        self._validate_chain_components(self.prompt, middle_chain)
-        
-        final = final_chain if final_chain is not None else self.output_parser
-        
-        if middle_chain:
-            func_chain = RunnableSequence(self.prompt, middle_chain, final)
-            return func_chain.invoke()
-        return None
-
-    def chain_parallel_invoke(self, parallel_chain: List) -> Any:
-        """
-        Chain invoke in parallel
-        Args:
-            parallel_chain: List of chains to run in parallel
-        Returns:
-            Any: Output from the parallel chains
-        Raises:
-            ImportError: If LangChain is not installed
-        """
-        if not check_package("langchain"):
-            raise ImportError("LangChain package not found. Please install it using: pip install langchain")
-        return parallel_chain.invoke()
-
-    def chain_branch_invoke(self, branch_chain: Dict) -> Any:
-        """
-        Chain invoke with branching
-        Args:
-            branch_chain: Dictionary of branch chains
-        Returns:
-            Any: Output from the branch chain
-        Raises:
-            ImportError: If LangChain is not installed
-        """
-        if not check_package("langchain"):
-            raise ImportError("LangChain package not found. Please install it using: pip install langchain")
-        return branch_chain.invoke()
-
-    @staticmethod
-    def create_parallel_chain(prompt_template: str, model: Any, branches: Dict[str, Any]) -> Any:
-        """
-        Create a parallel chain
-        Args:
-            prompt_template: Template for the prompt
-            model: The language model
-            branches: Dictionary of branch configurations
-        Returns:
-            Any: Configured parallel chain
-        """
-        from langchain.schema.runnable import RunnableParallel
-        return (
-            prompt_template
-            | model
-            | StrOutputParser()
-            | RunnableParallel(branches=branches)
-        )
-
-    @staticmethod
-    def create_branch_chain(conditions: List[tuple], default_chain: Any) -> Any:
-        """
-        Create a branch chain
-        Args:
-            conditions: List of condition-chain tuples
-            default_chain: Default chain to use
-        Returns:
-            Any: Configured branch chain
-        """
-        from langchain.schema.runnable import RunnableBranch
-        return RunnableBranch(*conditions, default_chain)
+## file for chaining functions in chatbot
+
+from typing import Optional, List, Dict, Any, Union
+from dataclasses import dataclass
+from langchain.schema.output_parser import StrOutputParser
+from mb_rag.chatbot.prompts import invoke_prompt
+from langchain.schema.runnable import RunnableLambda, RunnableSequence
+from mb_rag.utils.extra import check_package
+
+__all__ = ['Chain', 'ChainConfig']
+
+def check_langchain_dependencies() -> None:
+    """
+    Check if required LangChain packages are installed
+    Raises:
+        ImportError: If any required package is missing
+    """
+    if not check_package("langchain"):
+        raise ImportError("LangChain package not found. Please install it using: pip install langchain")
+    if not check_package("langchain_core"):
+        raise ImportError("LangChain Core package not found. Please install it using: pip install langchain-core")
+
+# Check dependencies before importing
+check_langchain_dependencies()
+
+@dataclass
+class ChainConfig:
+    """Configuration for chain operations"""
+    prompt: Optional[str] = None
+    prompt_template: Optional[str] = None
+    input_dict: Optional[Dict[str, Any]] = None
+
+class Chain:
+    """
+    Class to chain functions in chatbot with improved OOP design
+    """
+    def __init__(self, model: Any, config: Optional[ChainConfig] = None, **kwargs):
+        """
+        Initialize chain
+        Args:
+            model: The language model to use
+            config: Chain configuration
+            **kwargs: Additional arguments
+        """
+        self.model = model
+        self._output_parser = StrOutputParser()
+        self._initialize_config(config, **kwargs)
+
+    @classmethod
+    def from_template(cls, model: Any, template: str, input_dict: Dict[str, Any], **kwargs) -> 'Chain':
+        """
+        Create chain from template
+        Args:
+            model: The language model
+            template: Prompt template
+            input_dict: Input dictionary for template
+            **kwargs: Additional arguments
+        Returns:
+            Chain: New chain instance
+        """
+        config = ChainConfig(
+            prompt_template=template,
+            input_dict=input_dict
+        )
+        return cls(model, config, **kwargs)
+
+    @classmethod
+    def from_prompt(cls, model: Any, prompt: str, **kwargs) -> 'Chain':
+        """
+        Create chain from direct prompt
+        Args:
+            model: The language model
+            prompt: Direct prompt
+            **kwargs: Additional arguments
+        Returns:
+            Chain: New chain instance
+        """
+        config = ChainConfig(prompt=prompt)
+        return cls(model, config, **kwargs)
+
+    def _initialize_config(self, config: Optional[ChainConfig], **kwargs) -> None:
+        """Initialize chain configuration"""
+        if config:
+            self.input_dict = config.input_dict
+            if config.prompt_template:
+                self.prompt = invoke_prompt(config.prompt_template, self.input_dict)
+            else:
+                self.prompt = config.prompt
+        else:
+            self.input_dict = kwargs.get('input_dict')
+            if prompt_template := kwargs.get('prompt_template'):
+                self.prompt = invoke_prompt(prompt_template, self.input_dict)
+            else:
+                self.prompt = kwargs.get('prompt')
+
+    @property
+    def output_parser(self) -> StrOutputParser:
+        """Get the output parser"""
+        return self._output_parser
+
+    @staticmethod
+    def _validate_chain_components(prompt: Any, middle_chain: Optional[List] = None) -> None:
+        """
+        Validate chain components
+        Args:
+            prompt: The prompt to validate
+            middle_chain: Optional middle chain to validate
+        Raises:
+            ValueError: If validation fails
+        """
+        if prompt is None:
+            raise ValueError("Prompt is not provided")
+        if middle_chain is not None and not isinstance(middle_chain, list):
+            raise ValueError("middle_chain should be a list")
+
+    def invoke(self) -> Any:
+        """
+        Invoke the chain
+        Returns:
+            Any: Output from the chain
+        Raises:
+            Exception: If prompt is not provided
+        """
+        self._validate_chain_components(self.prompt)
+        chain_output = self.prompt | self.model | self.output_parser
+        return chain_output
+
+    def chain_sequence_invoke(self, 
+                            middle_chain: Optional[List] = None,
+                            final_chain: Optional[RunnableLambda] = None) -> Any:
+        """
+        Chain invoke the sequence
+        Args:
+            middle_chain: List of functions/Prompts/RunnableLambda to chain
+            final_chain: Final chain to run
+        Returns:
+            Any: Output from the chain
+        """
+        self._validate_chain_components(self.prompt, middle_chain)
+        
+        final = final_chain if final_chain is not None else self.output_parser
+        
+        if middle_chain:
+            func_chain = RunnableSequence(self.prompt, middle_chain, final)
+            return func_chain.invoke()
+        return None
+
+    def chain_parallel_invoke(self, parallel_chain: List) -> Any:
+        """
+        Chain invoke in parallel
+        Args:
+            parallel_chain: List of chains to run in parallel
+        Returns:
+            Any: Output from the parallel chains
+        Raises:
+            ImportError: If LangChain is not installed
+        """
+        if not check_package("langchain"):
+            raise ImportError("LangChain package not found. Please install it using: pip install langchain")
+        return parallel_chain.invoke()
+
+    def chain_branch_invoke(self, branch_chain: Dict) -> Any:
+        """
+        Chain invoke with branching
+        Args:
+            branch_chain: Dictionary of branch chains
+        Returns:
+            Any: Output from the branch chain
+        Raises:
+            ImportError: If LangChain is not installed
+        """
+        if not check_package("langchain"):
+            raise ImportError("LangChain package not found. Please install it using: pip install langchain")
+        return branch_chain.invoke()
+
+    @staticmethod
+    def create_parallel_chain(prompt_template: str, model: Any, branches: Dict[str, Any]) -> Any:
+        """
+        Create a parallel chain
+        Args:
+            prompt_template: Template for the prompt
+            model: The language model
+            branches: Dictionary of branch configurations
+        Returns:
+            Any: Configured parallel chain
+        """
+        from langchain.schema.runnable import RunnableParallel
+        return (
+            prompt_template
+            | model
+            | StrOutputParser()
+            | RunnableParallel(branches=branches)
+        )
+
+    @staticmethod
+    def create_branch_chain(conditions: List[tuple], default_chain: Any) -> Any:
+        """
+        Create a branch chain
+        Args:
+            conditions: List of condition-chain tuples
+            default_chain: Default chain to use
+        Returns:
+            Any: Configured branch chain
+        """
+        from langchain.schema.runnable import RunnableBranch
+        return RunnableBranch(*conditions, default_chain)

mb_rag/chatbot/conversation.py

@@ -0,0 +1,185 @@
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from typing import Optional, List, Any, Union
+
+__all__ = [
+    'ConversationModel'
+]
+
+class ConversationModel:
+    """
+    A class to handle conversation with AI models
+    
+    Attributes:
+        chatbot: The AI model for conversation
+        message_list (List): List of conversation messages
+        file_path (str): Path to save/load conversations. Can be local or S3
+    """
+    
+    def __init__(self, 
+                llm: Any,
+                message_list: Optional[List[Any]] = None,
+                file_path: Optional[str] = None,
+                **kwargs) -> None:
+        """Initialize conversation model"""
+        self.chatbot = llm
+        if message_list:
+            self.message_list = message_list
+        else:
+            self.message_list = []
+        if file_path:
+            self.file_path = file_path
+        else:
+            self.file_path = None
+
+    def initialize_conversation(self,context_message: str = "") -> None:
+        """Initialize conversation state.
+            Getting the content from file_path if provided"""
+        if self.file_path:
+            self.load_conversation()
+
+        if context_message:
+            self.message_list.append(SystemMessage(content=context_message))
+        else:
+            self.message_list.append(SystemMessage(content="""This is conversation model. 
+                            Look into the conversation history and answer the question if provided.
+                            Give a brief introduction of the conversation history."""))
+        message_list_content = "".join(self.all_messages_content)
+        return self.add_message(message_list_content,get_content_only=True)    
+
+    def _ask_question(self,query: str,images: list = None, 
+                     get_content_only: bool = True) -> str:
+        """
+        Ask a question and get response
+        Args:
+            query: Question to ask
+            get_content_only: Whether to return only content
+        Returns:
+            str: Response from the model
+        """
+        if images:
+            res = self.chatbot.invoke_query(query,images=images,get_content_only=get_content_only)
+        else:
+            res = self.chatbot.invoke_query(query,get_content_only=get_content_only)
+        return res
+
+    def add_message(self, query: str,images: list = None,get_content_only: bool = True) -> str:
+        """
+        Add a message to the conversation
+        Args:
+            query (str): Question to ask
+            images (list): List of images to send to the model
+            get_content_only (bool): Whether to return only content
+        Returns:
+            str: Response from the chatbot
+        """
+        self.message_list.append(HumanMessage(content=query))
+        res = self._ask_question(query,images=images,get_content_only=get_content_only)
+        self.message_list.append(AIMessage(content=res))
+        return res
+
+    @property
+    def all_messages(self) -> List[Union[SystemMessage, HumanMessage, AIMessage]]:
+        """Get all messages"""
+        return self.message_list
+
+    @property
+    def last_message(self) -> str:
+        """Get the last message"""
+        return self.message_list[-1].content
+
+    @property
+    def all_messages_content(self) -> List[str]:
+        """Get content of all messages"""
+        return [message.content for message in self.message_list]
+
+    def _is_s3_path(self, path: str) -> bool:
+        """
+        Check if path is an S3 path
+        Args:
+            path (str): Path to check
+        Returns:
+            bool: True if S3 path
+        """
+        return path.startswith("s3://")
+
+    def save_conversation(self, file_path: Optional[str] = None, **kwargs) -> bool:
+        """
+        Save the conversation
+        Args:
+            file_path: Path to save the conversation
+            **kwargs: Additional arguments for S3
+        Returns:
+            bool: Success status
+        """
+        if self._is_s3_path(file_path or self.file_path):
+            print("Saving conversation to S3.")
+            self.save_file_path = file_path
+            return self._save_to_s3(self.file_path,**kwargs)
+        return self._save_to_file(file_path or self.file_path)
+
+    def _save_to_s3(self,**kwargs) -> bool:
+        """Save conversation to S3"""
+        try:
+            client = kwargs.get('client', self.client)
+            bucket = kwargs.get('bucket', self.bucket)
+            client.put_object(
+                Body=str(self.message_list),
+                Bucket=bucket,
+                Key=self.save_file_path
+            )
+            print(f"Conversation saved to s3_path: {self.s3_path}")
+            return True
+        except Exception as e:
+            raise ValueError(f"Error saving conversation to s3: {e}")
+
+    def _save_to_file(self, file_path: str) -> bool:
+        """Save conversation to file"""
+        try:
+            with open(file_path, 'w') as f:
+                for message in self.message_list:
+                    f.write(f"{message.content}\n")
+            print(f"Conversation saved to file: {file_path}")
+            return True
+        except Exception as e:
+            raise ValueError(f"Error saving conversation to file: {e}")
+
+    def load_conversation(self, file_path: Optional[str] = None, **kwargs) -> List[Any]:
+        """
+        Load a conversation
+        Args:
+            file_path: Path to load from
+            **kwargs: Additional arguments for S3
+        Returns:
+            List: Loaded messages
+        """
+        self.message_list = []
+        if self._is_s3_path(file_path or self.file_path):
+            print("Loading conversation from S3.")
+            self.file_path = file_path
+            return self._load_from_s3(**kwargs)
+        return self._load_from_file(file_path or self.file_path)
+
+    def _load_from_s3(self, **kwargs) -> List[Any]:
+        """Load conversation from S3"""
+        try:
+            client = kwargs.get('client', self.client)
+            bucket = kwargs.get('bucket', self.bucket)
+            res = client.get_response(client, bucket, self.s3_path)
+            res_str = eval(res['Body'].read().decode('utf-8'))
+            self.message_list = [SystemMessage(content=res_str)]
+            print(f"Conversation loaded from s3_path: {self.file_path}")
+            return self.message_list
+        except Exception as e:
+            raise ValueError(f"Error loading conversation from s3: {e}")
+
+    def _load_from_file(self, file_path: str) -> List[Any]:
+        """Load conversation from file"""            
+        try:
+            with open(file_path, 'r') as f:
+                lines = f.readlines()
+            for line in lines:
+                self.message_list.append(SystemMessage(content=line))
+            print(f"Conversation loaded from file: {file_path}")
+            return self.message_list
+        except Exception as e:
+            raise ValueError(f"Error loading conversation from file: {e}")