nick1221/system1

1import os2import datetime3import json4import re5import logging6import pprint7import sys8import uuid9import argparse10from dotenv import load_dotenv11from langchain import hub12from langchain.globals import set_debug13from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder14from langchain_core.messages import HumanMessage, AIMessage, SystemMessage15from llm_providers import get_openai_llm, get_bedrock_llm, get_bedrock_chat_llm, get_together_llm, get_groq_llm, get_ollama_llm16from dotenv import load_dotenv17load_dotenv()18 19AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY')20AWS_SECRET_ACCESS_KEY = os.getenv('AWS_SECRET_ACCESS_KEY')21AWS_REGION = os.getenv('AWS_REGION', 'us-east-1')22 23from langgraph.prebuilt import create_react_agent24from langgraph.checkpoint import MemorySaver25 26class System1:27    def __init__(self, llm_provider='openai', model_id='gpt-4-1106-preview', debug=False, agent_type="structured", toolkits=['jira', 'launch', 'time']):28        self.model_id = model_id29        self.agent_type = agent_type30        self.tools = self.load_tools(toolkits)31        self._setup_logging(debug)32        self.llm = self._get_llm(llm_provider)33        self.memory = MemorySaver()34        self.chat_history = []35        self.agent = create_react_agent(self.llm, self.tools, checkpointer=self.memory, messages_modifier=self._get_system_prompt())36 37    def _setup_logging(self, debug):38        self.log_dir = "output/logs"39        os.makedirs(self.log_dir, exist_ok=True)40        timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")41        self.log_file = os.path.join(self.log_dir, f"vector_log_{timestamp}.txt")42        if debug:43            set_debug(True)44 45 46    def _get_llm(self, llm_provider):47        llm_providers = {48            'openai': get_openai_llm,49            'bedrock': get_bedrock_llm,50            'bedrock-chat': get_bedrock_chat_llm,51            'together': get_together_llm,52            'groq': get_groq_llm,53            'ollama': get_ollama_llm54        }55 56        if llm_provider in llm_providers:57            return llm_providers[llm_provider](self.model_id)58        else:59            raise ValueError(f"Unsupported LLM provider: {llm_provider}")60        61    def _get_system_prompt(self):62        return """You are a very powerful assistant named Vector63        Your goal is to accomplish business objectives with the tools at your disposal. 64        You must try to accomplish every task, no matter how difficult.65        DO NOT MENTION TOOL NAMES IN YOUR FINAL RESPONSE!!!66        Provide your final response in a well formatted human readable style. DO NOT OUTPUT JSON IN FINAL RESPONSE. ONLY USE JSON FOR TOOL CALLING67        Ignore null or none values unless the user specifically requests that information68        Provide all relevant information pertaining to the user's request in your final output69        """70 71    async def chat_stream(self, chat_string):72        try:73            print(chat_string)74            inputs = {"input": str(chat_string), "chat_history": self.chat_history}75            76            async for chunk in self.agent_executor.astream(inputs):77                yield str(chunk)78                print("------")79                pprint.pprint(chunk, depth=1)80        except Exception as e:81            yield str(e)82 83    def chat_messages(self, messages):84        config = {"configurable": {"thread_id": uuid.uuid1()}}85        response = self.agent.invoke({"messages": messages}, config=config)86        return response87 88    def chat_llm(self, chat_input: str):89        return self.llm.invoke([HumanMessage(content=chat_input)])90        return self.llm.invoke(chat_string)91 92    def load_tools(self, toolkits):93        workdir = os.path.join(os.getcwd(), 'tmp')94        from tools import AgentTools95        return AgentTools(toolkits=toolkits, workdir=workdir).get_tools()96 97 98if __name__ == '__main__':99    logging.basicConfig(stream=sys.stdout, level=logging.INFO)100    logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))101    load_dotenv()102 103    parser = argparse.ArgumentParser(description='Vector: A conversational AI agent')104    parser.add_argument('--chat', action='store_true', help='Run Vector in chat mode')105    parser.add_argument('--query', type=str, help='Run Vector in query mode with a single input')106    parser.add_argument('--local', action='store_true', help='Run Vector in local mode')107    parser.add_argument('--llm', action='store_true', help='Run Vector in local mode')108 109    args = parser.parse_args()110 111    if args.chat:112        #vector = Vector(llm_provider="bedrock-chat", model_id="anthropic.claude-3-haiku-20240307-v1:0", agent_type="react", toolkits=['jira'])113        #vector = Vector(llm_provider="ollama", model_id="phi3")114        vector = System1(llm_provider="openai", model_id="gpt-4o", agent_type="react", toolkits=['jira'])115 116        while True:117            user_input = input('You > ')118            if user_input.lower() == 'quit':119                break120            else:121                messages = [("user", user_input)]122                result = vector.chat_messages(messages)123                print()124 125    if args.local:126        vector = System1(llm_provider="ollama", model_id="phi3:latest", toolkits=["launch"])127        from langchain_core.messages.human import HumanMessage128        messages = []129        while True:130            user_input = input('You > ')131            if user_input.lower() == 'quit':132                break133            else:134                result = vector.chat_messages([(HumanMessage(content=user_input))])135                print(result)136        137    if args.llm:138        vector = System1(llm_provider="bedrock-chat", model_id="anthropic.claude-3-haiku-20240307-v1:0")139        result = vector.chat_llm(input("> "))140        print("Vector LLM: " + result.content)141 142