Upload 6 files

rmrkl/__init__.py

@@ -0,0 +1,3 @@
+from .version import __version__
+from .agent import ChatZeroShotAgent
+from .executor import RetryAgentExecutor

rmrkl/agent.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+from langchain.chat_models.base import BaseChatModel
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from langchain.callbacks.base import BaseCallbackManager
+from .prompts import FORMAT_INSTRUCTIONS, SUFFIX, QUESTION_PROMPT, PREFIX
+from langchain.agents.agent import Agent, AgentOutputParser
+from typing import Any, Optional, Sequence
+from langchain.tools import BaseTool
+from langchain.agents.mrkl.base import ZeroShotAgent
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+    AIMessagePromptTemplate,
+)
+from .output_parser import ChatZeroShotOutputParser
+
+
+class ChatZeroShotAgent(ZeroShotAgent):
+    """Agent for the MRKL chain."""
+
+    @classmethod
+    def create_prompt(
+        cls,
+        tools: Sequence[BaseTool],
+        prefix: str = PREFIX,
+        suffix: str = SUFFIX,
+        format_instructions: str = FORMAT_INSTRUCTIONS,
+        question_prompt: str = QUESTION_PROMPT,
+    ) -> PromptTemplate:
+        """Create prompt in the style of the zero shot agent.
+
+        Args:
+            tools: List of tools the agent will have access to, used to format the
+                prompt.
+            prefix: String to put before the list of tools.
+            suffix: String to put after the list of tools.
+            input_variables: List of input variables the final prompt will expect.
+
+        Returns:
+            A PromptTemplate with the template assembled from the pieces here.
+        """
+        tool_strings = "\n".join(
+            [f"    {tool.name}: {tool.description}" for tool in tools]
+        )
+        tool_names = ", ".join([tool.name for tool in tools])
+        format_instructions = format_instructions.format(
+            tool_names=tool_names, tool_strings=tool_strings
+        )
+        human_prompt = PromptTemplate(
+            template=question_prompt,
+            input_variables=["input"],
+            partial_variables={"tool_strings": tool_strings},
+        )
+        human_message_prompt = HumanMessagePromptTemplate(prompt=human_prompt)
+        ai_message_prompt = AIMessagePromptTemplate.from_template(suffix)
+        system_message_prompt = SystemMessagePromptTemplate.from_template(
+            '\n\n'.join(
+                [
+                    prefix,
+                    format_instructions
+                ]
+            )
+        )
+        # ignore suffix
+        return ChatPromptTemplate.from_messages(
+            [system_message_prompt, human_message_prompt, ai_message_prompt]
+        )
+
+    @classmethod
+    def from_llm_and_tools(
+        cls,
+        llm: BaseChatModel,
+        tools: Sequence[BaseTool],
+        callback_manager: Optional[BaseCallbackManager] = None,
+        output_parser: Optional[AgentOutputParser] = ChatZeroShotOutputParser(),
+        prefix: str = PREFIX,
+        suffix: str = SUFFIX,
+        format_instructions: str = FORMAT_INSTRUCTIONS,
+        question_prompt: str = QUESTION_PROMPT,
+        **kwargs: Any,
+    ) -> Agent:
+        """Construct an agent from an LLM and tools."""
+        cls._validate_tools(tools)
+        prompt = cls.create_prompt(
+            tools,
+            prefix=prefix,
+            suffix=suffix,
+            format_instructions=format_instructions,
+            question_prompt=question_prompt,
+        )
+        llm_chain = LLMChain(
+            llm=llm,
+            prompt=prompt,
+            callback_manager=callback_manager,
+        )
+        tool_names = [tool.name for tool in tools]
+        _output_parser = output_parser or cls._get_default_output_parser()
+        return cls(
+            llm_chain=llm_chain,
+            allowed_tools=tool_names,
+            output_parser=_output_parser,
+            **kwargs,
+        )

rmrkl/executor.py

@@ -0,0 +1,23 @@
+from typing import Dict, List, Tuple, Union, Optional
+
+from langchain.agents import AgentExecutor
+from langchain.schema import AgentAction, AgentFinish, OutputParserException
+from langchain.tools import BaseTool
+from langchain.callbacks.manager import CallbackManagerForChainRun
+
+
+class ExceptionTool(BaseTool):
+    name: str = "_Exception"
+    description: str = "Exception tool"
+
+    def _run(self, query: str)  :
+        return query
+
+    async def _arun(self, query: str) :
+        return query
+
+
+class RetryAgentExecutor(AgentExecutor):
+    """Agent executor that retries on output parser exceptions."""
+    # for backwards compatibility
+    handle_parsing_errors: bool = True

rmrkl/output_parser.py

@@ -0,0 +1,37 @@
+import re
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+import langchain
+from langchain import LLMChain
+from langchain.agents.agent import AgentOutputParser
+from langchain.schema import AgentAction, AgentFinish, OutputParserException
+
+from .prompts import (FINAL_ANSWER_ACTION, FORMAT_INSTRUCTIONS,
+                      QUESTION_PROMPT, SUFFIX)
+
+
+class ChatZeroShotOutputParser(AgentOutputParser):
+    def get_format_instructions(self) -> str:
+        return FORMAT_INSTRUCTIONS
+
+    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
+        if FINAL_ANSWER_ACTION in text:
+            return AgentFinish(
+                {"output": text.split(FINAL_ANSWER_ACTION)[-1].strip()}, text
+            )
+
+
+        # Remove 'Thought' SUFFIX
+        if text.startswith('Thought:'):
+            text = text[8:]
+
+        # \s matches against tab/newline/whitespace
+        regex = (
+            r"Action\s*\d*\s*:[\s]*(.*?)[\s]*Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)"
+        )
+        match = re.search(regex, text, re.DOTALL)
+        if not match:
+            raise OutputParserException(f"Could not parse LLM output: `{text}`")
+        action = match.group(1).strip()
+        action_input = match.group(2)
+        return AgentAction(action, action_input.strip(" ").strip('"'), text.strip())

rmrkl/prompts.py

@@ -0,0 +1,37 @@
+# flake8: noqa
+
+PREFIX = """
+You are an AI system called TeLLAgent and your task is to respond to the question or
+solve the problem to the best of your ability using the provided tools.
+
+""" 
+
+FORMAT_INSTRUCTIONS = """
+You should only respond with a single complete
+Thought, Action, Action Input format
+OR a single Final Answer format.
+
+Complete Format:
+
+Thought: (reflect on your progress and decide what to do next)
+Action: (the action name, should be one of [{tool_names}])
+Action Input: (the input string to the action)
+
+OR
+
+Final Answer: (the final answer to the original input question)
+
+"""
+QUESTION_PROMPT = """
+Answer the question below using the following tools:
+Note: Direct output the results when using the following tools:Imageanalysis, codewriter, pdfreader, rag  without further processing.
+
+{tool_strings}
+
+Question: {input}
+"""
+SUFFIX = """
+Direct output the results when using Imageanalysis, codewriter, pdfreader, rag tools without further processing.
+Thought: {agent_scratchpad}
+"""
+FINAL_ANSWER_ACTION = "Final Answer:"

rmrkl/version.py

@@ -0,0 +1 @@
+__version__ = "0.0.3"