ZurichAgent - LlamaIndex Prototype

zurich_agent.py

class ZurichAgent:
    def __init__(self, wordlift_key: str, openai_key: str, openapi_yaml: str, system_prompt: str, chat_memory: ChatMemoryBuffer):
        # read sys prompt
        with open(system_prompt, 'r') as f:
            prompt_text = f.read()

        # read openapi
        with open(openapi_yaml, 'r') as f:
            openapi_spec = OpenAPIToolSpec(
                yaml.safe_load(f)
            )
        requests_spec = RequestsToolSpec({
            "api.wordlift.io": {
                "Authorization": f"Key {wordlift_key}",
                "Content-Type": "application/json",
            }
        })

        self._agent = OpenAIAgent.from_tools(
            [*openapi_spec.to_tool_list(), *requests_spec.to_tool_list()],
            llm=OpenAI(model="gpt-4o-mini", api_key=openai_key),
            system_prompt=prompt_text,
            verbose=True,
            memory=chat_memory,
        )

    def handle(self, message: str) -> str:
        return self._agent.chat(message).response

def create_zurich_agent():
    """
    Create and initialize a ZurichAgent instance using environment variables.
    """
    # Initialize chat memory
    chat_memory = ChatMemoryBuffer.from_defaults(token_limit=4000)

    # Create agent instance
    agent = ZurichAgent(
        wordlift_key=WORDLIFT_API_KEY,
        openai_key=OPENAI_API_KEY,
        openapi_yaml='/content/openapi.yaml',
        system_prompt='/content/prompt.md',  # Pass prompt content
        chat_memory=chat_memory
    )

    return agent