Interfaces

from typing import Any, Dict, List, Optional
from fairo.core.chat.chat import ChatFairo
from langchain.tools import tool
from pydantic import BaseModel, Field
from langchain.agents import AgentExecutor, create_json_chat_agent
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from fairo.core.workflow.dependency import FairoVectorStore
from langchain_core.messages import BaseMessage
from fairo.core.utils import parse_chat_interface_output

class InputSchema(BaseModel):
    input: str = Field(description="The user input question")
    chat_history: List[BaseMessage] = Field(
        description="List of previous conversation messages"
    )

class FaqAgent:
    prompt = """
You are a helpful support agent for Best CRM Inc.

You MUST always follow this format:
- You must remember user data provided by the user
- If you need to look something up, output a JSON with the key "action" set to one of the provided tools, and "action_input" with the query.
- If you are ready to answer, output a JSON with the key "action" set to "Final Answer" and the key "action_input" with your answer text.
- You may ONLY use the provided tool names. Do not invent new ones.

Available tools: {tools}
Tool names: {tool_names}
""".strip()
    llm = ChatFairo()

    class FaqQueryArgs(BaseModel):
        query: str = Field(
            ...,
            description="Query to search for information about Best CRM Inc., including FAQs, policies, and setup instructions",
        )

    @tool(args_schema=FaqQueryArgs)
    def search_faq_knowledge_store(query: str) -> str:
        """
            Search the Best CRM Inc. knowledge base for information about the CRM software,
            policies, setup instructions, and frequently asked questions.
        """
        knowledge_store = FairoVectorStore(
            collection_name="faq",
        )
        docs = knowledge_store.similarity_search(query, k=5)
        if not docs:
            return "I don't have information about that topic in my knowledge base."
        return "\n\n".join(doc.page_content for doc in docs)
    

    tools = [search_faq_knowledge_store]

    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system", prompt
            ),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            ("ai", "{agent_scratchpad}"),
        ]
    )

    def invoke(self, inputs: Dict[str, Any], config: Optional[Any] = None,) -> Dict[str, Any]:
        agent = create_json_chat_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=self.prompt,
        )

        executor = AgentExecutor(
            name="FAQ Chat Agent",
            agent=agent,
            tools=self.tools,
            verbose=True,
            handle_parsing_errors=True,
            max_iterations=5,
            input_schema=InputSchema,
            return_intermediate_steps=True,
        )
        result = executor.invoke(inputs, config=config)

        messages = parse_chat_interface_output(result)
        return messages

def faq_agent():
    """
        Plot generation assistant
    """
    return FaqAgent()

from fairo.core.execution.executor import FairoExecutor
from agent import faq_agent, InputSchema

if __name__ == "__main__":
    result = FairoExecutor(
        input_schema=InputSchema,
        agents=[faq_agent],
    ).run({"input": "use the plot_html_chart tool", "chat_history": []})