flo-ai 0.0.1__tar.gz

flo_ai-0.0.1/PKG-INFO

@@ -0,0 +1,101 @@
+Metadata-Version: 2.1
+Name: flo-ai
+Version: 0.0.1
+Summary: A easy way to create structured AI agents
+License: MIT
+Author: vizsatiz
+Author-email: vishnu@rootfor.xyz
+Requires-Python: >=3.10,<3.13
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: google-cloud-bigquery (>=3.23.1,<4.0.0)
+Requires-Dist: httpx (==0.27.0)
+Requires-Dist: ipython (==8.24.0)
+Requires-Dist: langchain (==0.2.1)
+Requires-Dist: langchain-chroma (>=0.1.1,<0.2.0)
+Requires-Dist: langchain-community (==0.2.0)
+Requires-Dist: langchain-experimental (==0.0.59)
+Requires-Dist: langchain-mistralai (>=0.1.7,<0.2.0)
+Requires-Dist: langchain-mongodb (>=0.1.5,<0.2.0)
+Requires-Dist: langchain-openai (>=0.1.8,<0.2.0)
+Requires-Dist: langgraph (==0.0.55)
+Requires-Dist: openai (==1.30.5)
+Requires-Dist: pillow (>=10.3.0,<11.0.0)
+Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
+Description-Content-Type: text/markdown
+
+# Flo: Composable AI Agents
+
+Flo gives you a composable framework for creating agentic AI architectures. What we intent to do here is to create an easy framework for AI app developers to compose apps using pre-implemented architectural components, at the same time providing the flexibilty to create their own components.
+
+### What is composibility ?
+
+Composility is the ability to use smaller components to build bigger apps. In a composible architecture, you will be given smaller building blocks which you can use to build a bigger applications. It is every similar to how legos work, you are given smaller lego blocks which when put together creates a whole structure.
+
+### What are the building blocks in Flo ?
+
+In flo, we tried to put togther a system where we have small micro components like vector store, or simple LLM prompts, and then higher components/arhictectures made of these mico components like RAGs, Agentic Teams etc.
+
+# Flo vs langraph/langchain
+
+Flo is built with langraph under the stood. So everything that works in langraph still works here, including all tools and architectures. The following makes flo easy to implement:
+
+1. Langraph needs good understanding of underlying graph and states, its more of a raw tool and asks developers to implement the required artitectures. While flo is more usecase friendly and components can be easily used by using the flo classes, which has lot of internal abstraction for ease of use.
+
+2. In every AI component that has gone in production their are lot of nuances that needs to be implemented to get production quality output, flo inherently implements the best architecutures and gives then out of the box, you can enable and disable as you wish. This not only reduces the complexity but also improves the time to iterate solutions
+
+3. Every component in flo is combosable meaning you can easily put them together and flo takes care of routing between the components where as in langraph the developer has to tie these up.
+
+# How to use
+
+Flo supports two ways to set up and run the components, first is through code. This is much more flexible. This can help you write your own tools and add them to the flo.
+
+Second way it to use yaml. You can write an yaml to define your agentic workflow or RAG and it compiles into a application. See examples below
+
+## Building your first agent by structured yaml
+To create a small team of researcher + blogger for writing blogs
+
+```python
+# This yaml defines a team of 2 agents + 1 supervisor
+yaml_data = """
+apiVersion: flo/alpha-v1
+kind: FloRoutedTeam
+name: blogging-team
+team:
+    name: BloggingTeam
+    router:
+        name: TeamLead
+        kind: supervisor
+    agents:
+      - name: Researcher
+        prompt: Do a research on the internet and find articles of relevent to the topic asked by the user, always try to find the latest information on the same
+        tools:
+        - name: TavilySearchResults
+      - name: Blogger
+        prompt: From the documents provider by the researcher write a blog of 300 words with can be readily published, make in engaging and add reference links to original blogs
+        tools:
+        - name: TavilySearchResults
+"""
+
+input_prompt = """
+Question: Write me an interesting blog about latest advancements in agentic AI
+"""
+
+llm = ChatOpenAI(temperature=0, model_name='gpt-4o')
+
+# Register all the tools at some place and use everywhere in the yaml
+session = FloSession(llm).register_tool(
+    name="TavilySearchResults", 
+    tool=TavilySearchResults()
+)
+
+# Build the final flow and run it
+flo: Flo = Flo.build_with_yaml(session, yaml=yaml_data)
+
+# call invoke or stream
+for s in flo.stream(input_prompt)
+```
+

flo_ai-0.0.1/README.md

@@ -0,0 +1,71 @@
+# Flo: Composable AI Agents
+
+Flo gives you a composable framework for creating agentic AI architectures. What we intent to do here is to create an easy framework for AI app developers to compose apps using pre-implemented architectural components, at the same time providing the flexibilty to create their own components.
+
+### What is composibility ?
+
+Composility is the ability to use smaller components to build bigger apps. In a composible architecture, you will be given smaller building blocks which you can use to build a bigger applications. It is every similar to how legos work, you are given smaller lego blocks which when put together creates a whole structure.
+
+### What are the building blocks in Flo ?
+
+In flo, we tried to put togther a system where we have small micro components like vector store, or simple LLM prompts, and then higher components/arhictectures made of these mico components like RAGs, Agentic Teams etc.
+
+# Flo vs langraph/langchain
+
+Flo is built with langraph under the stood. So everything that works in langraph still works here, including all tools and architectures. The following makes flo easy to implement:
+
+1. Langraph needs good understanding of underlying graph and states, its more of a raw tool and asks developers to implement the required artitectures. While flo is more usecase friendly and components can be easily used by using the flo classes, which has lot of internal abstraction for ease of use.
+
+2. In every AI component that has gone in production their are lot of nuances that needs to be implemented to get production quality output, flo inherently implements the best architecutures and gives then out of the box, you can enable and disable as you wish. This not only reduces the complexity but also improves the time to iterate solutions
+
+3. Every component in flo is combosable meaning you can easily put them together and flo takes care of routing between the components where as in langraph the developer has to tie these up.
+
+# How to use
+
+Flo supports two ways to set up and run the components, first is through code. This is much more flexible. This can help you write your own tools and add them to the flo.
+
+Second way it to use yaml. You can write an yaml to define your agentic workflow or RAG and it compiles into a application. See examples below
+
+## Building your first agent by structured yaml
+To create a small team of researcher + blogger for writing blogs
+
+```python
+# This yaml defines a team of 2 agents + 1 supervisor
+yaml_data = """
+apiVersion: flo/alpha-v1
+kind: FloRoutedTeam
+name: blogging-team
+team:
+    name: BloggingTeam
+    router:
+        name: TeamLead
+        kind: supervisor
+    agents:
+      - name: Researcher
+        prompt: Do a research on the internet and find articles of relevent to the topic asked by the user, always try to find the latest information on the same
+        tools:
+        - name: TavilySearchResults
+      - name: Blogger
+        prompt: From the documents provider by the researcher write a blog of 300 words with can be readily published, make in engaging and add reference links to original blogs
+        tools:
+        - name: TavilySearchResults
+"""
+
+input_prompt = """
+Question: Write me an interesting blog about latest advancements in agentic AI
+"""
+
+llm = ChatOpenAI(temperature=0, model_name='gpt-4o')
+
+# Register all the tools at some place and use everywhere in the yaml
+session = FloSession(llm).register_tool(
+    name="TavilySearchResults", 
+    tool=TavilySearchResults()
+)
+
+# Build the final flow and run it
+flo: Flo = Flo.build_with_yaml(session, yaml=yaml_data)
+
+# call invoke or stream
+for s in flo.stream(input_prompt)
+```

flo_ai-0.0.1/flo_ai/__init__.py

@@ -0,0 +1,5 @@
+from flo.core import Flo
+from flo.models.flo_agent import FloAgentBuilder
+from flo.models.flo_supervisor import FloSupervisorBuilder
+from flo.models.flo_team import FloTeamBuilder
+from flo.state.flo_session import FloSession

flo_ai-0.0.1/flo_ai/builders/yaml_builder.py

@@ -0,0 +1,69 @@
+from flo.models.flo_team import FloTeamBuilder
+from flo.models.flo_agent import FloAgentBuilder, FloAgent
+from flo.models.flo_supervisor import FloSupervisorBuilder
+from flo.yaml.flo_team_builder import (FloRoutedTeamConfig, TeamConfig,
+                                        AgentConfig, FloAgentConfig)
+from flo.models.flo_executable import ExecutableFlo
+from flo.models.flo_planner import FloPlannerBuilder
+from flo.state.flo_session import FloSession
+from typing import Union, List
+
+def build_supervised_team(
+        session: FloSession,
+        flo_config: Union[FloRoutedTeamConfig, FloAgentConfig]) -> ExecutableFlo:
+    if isinstance(flo_config, FloRoutedTeamConfig):
+        team_config: TeamConfig = flo_config.team
+        team = parse_and_build_subteams(session, team_config, session.tools)    
+        return team
+    elif isinstance(flo_config, FloAgentConfig):
+        agent_config: AgentConfig = flo_config.agent
+        agent = create_agent(session, agent_config, session.tools)
+        return agent
+
+def parse_and_build_subteams(
+        session: FloSession,
+        team: TeamConfig,
+        tool_map) -> ExecutableFlo:
+    flo_team = None
+    if team.agents:
+        agents = []
+        for agent in team.agents:
+            flo_agent: FloAgent = create_agent(session, agent, tool_map)
+            agents.append(flo_agent)
+        router = create_router(session, team, agents)
+        flo_team = FloTeamBuilder(
+            session=session,
+            name=team.name,
+            supervisor=router
+        ).build()
+        if team.planner is not None:
+            return FloPlannerBuilder(session, team.planner.name, flo_team).build()
+    else:
+        flo_teams = []
+        for subteam in team.subteams:
+            flo_subteam = parse_and_build_subteams(session, subteam, tool_map)
+            flo_teams.append(flo_subteam)
+        router = create_router(session, team, agents)
+        flo_team = FloTeamBuilder(
+            session=session,
+            name=team.name,
+            supervisor=router
+        ).build()
+    return flo_team
+
+def create_router(session: FloSession, team_config: TeamConfig, agents: List[FloAgent]):
+    if (team_config.router.kind == "supervisor"):
+            return FloSupervisorBuilder(session, team_config.router.name, agents).build()
+    else:
+        raise Exception("Unknown router type")
+        
+def create_agent(session: FloSession, agent: AgentConfig, tool_map) -> FloAgent:
+    tools = [tool_map[tool.name] for tool in agent.tools]
+    flo_agent: FloAgent = FloAgentBuilder(
+        session,
+        agent.name, 
+        agent.prompt,
+        tools
+    ).build()
+    return flo_agent
+

flo_ai-0.0.1/flo_ai/constants/prompt_constants.py

@@ -0,0 +1 @@
+FLO_FINISH = "FINISH"

flo_ai-0.0.1/flo_ai/core.py

@@ -0,0 +1,40 @@
+from flo.yaml.flo_team_builder import to_supervised_team
+from flo.builders.yaml_builder import build_supervised_team, FloRoutedTeamConfig
+from typing import (
+    Any,
+    Iterator,
+    Union
+)
+from flo.state.flo_session import FloSession
+from flo.models.flo_executable import ExecutableFlo
+
+class Flo:
+
+    def __init__(self,
+                 session: FloSession,
+                 config: FloRoutedTeamConfig) -> None:
+        self.config = config
+        self.runnable: ExecutableFlo = build_supervised_team(session, config)
+
+    def stream(self, query, config = None) -> Iterator[Union[dict[str, Any], Any]]:
+        return self.runnable.stream(query, config)
+    
+    def invoke(self, query, config = None) -> Iterator[Union[dict[str, Any], Any]]:
+        return self.runnable.invoke(query, config)
+    
+    @staticmethod
+    def build_with_yaml(session: FloSession, yaml: str):
+        return Flo(session, to_supervised_team(yaml))
+
+    def draw(self, xray=True):
+        return self.runnable.draw(xray)
+    
+    def draw_to_file(self, filename: str, xray=True):
+        from PIL import Image as PILImage
+        import io
+        byte_image = self.runnable.draw(xray)
+        with io.BytesIO(byte_image) as image_io:
+            image = PILImage.open(image_io)
+            image.save(filename)
+
+        

flo_ai-0.0.1/flo_ai/helpers/utils.py

@@ -0,0 +1,10 @@
+import random
+import string
+
+def random_str(length: int = 5):
+    letters = string.ascii_letters + string.digits
+    result_str = ''.join(random.choice(letters) for i in range(length))
+    return result_str
+
+def randomize_name(name: str):
+    return "{}-{}".format(name, random_str(5))

flo_ai-0.0.1/flo_ai/models/exception.py

@@ -0,0 +1,7 @@
+class FloValidationException(Exception):
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)
+
+class FloIllegalStateException(Exception):
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)

flo_ai-0.0.1/flo_ai/models/flo_agent.py

@@ -0,0 +1,57 @@
+from langchain_core.tools import BaseTool
+from langchain.agents import AgentExecutor
+from langchain.agents import create_openai_tools_agent
+from langchain_core.runnables import Runnable
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from flo.helpers.utils import randomize_name
+from flo.models.flo_executable import ExecutableFlo
+from flo.state.flo_session import FloSession
+from typing import Union
+
+class FloAgent(ExecutableFlo):
+    def __init__(self, 
+                 agent: Runnable, 
+                 executor: AgentExecutor, 
+                 name: str) -> None:
+        super().__init__(name, executor, "agent")
+        self.agent: Runnable =  agent,
+        self.executor: AgentExecutor = executor
+
+class FloAgentBuilder:
+    def __init__(self, 
+                 session: FloSession,
+                 name: str, 
+                 prompt: Union[ChatPromptTemplate, str], 
+                 tools: list[BaseTool],
+                 verbose: bool = True,
+                 llm: Union[BaseLanguageModel, None] =  None,
+                 return_intermediate_steps: bool = False,
+                 handle_parsing_errors: bool = True) -> None:
+        self.name: str = randomize_name(name)
+        self.llm = llm if llm is not None else session.llm
+        # TODO improve to add more context of what other agents are available
+        self.prompt: ChatPromptTemplate = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    prompt,
+                ),
+                MessagesPlaceholder(variable_name="messages"),
+                MessagesPlaceholder(variable_name="agent_scratchpad")
+            ]
+        ) if isinstance(prompt, str) else prompt
+        self.tools: list[BaseTool] = tools
+        self.verbose = verbose
+        self.return_intermediate_steps = return_intermediate_steps
+        self.handle_parsing_errors = handle_parsing_errors
+
+
+    def build(self) -> AgentExecutor:
+        agent = create_openai_tools_agent(self.llm, self.tools, self.prompt)
+        executor = AgentExecutor(agent=agent, 
+                             tools=self.tools, 
+                             verbose=self.verbose, 
+                             return_intermediate_steps=True, 
+                             handle_parsing_errors=True)
+        return FloAgent(agent, executor, self.name)

flo_ai-0.0.1/flo_ai/models/flo_executable.py

@@ -0,0 +1,22 @@
+from flo.models.flo_member import FloMember
+from langgraph.graph.graph import CompiledGraph
+from langchain_core.messages import HumanMessage
+
+class ExecutableFlo(FloMember):
+    def __init__(self, name: str, graph: CompiledGraph, type: str = "team") -> None:
+        super().__init__(name, type)
+        self.graph = graph
+
+    def stream(self, work, config = None):
+        return self.graph.stream({
+             "messages": [
+                HumanMessage(content=work)
+            ]
+        }, config)
+    
+    def invoke(self, work):
+        return self.graph.invoke({
+             "messages": [
+                HumanMessage(content=work)
+            ]
+        })

flo_ai-0.0.1/flo_ai/models/flo_member.py

@@ -0,0 +1,4 @@
+class FloMember():
+    def __init__(self, name: str, type: str) -> None:
+        self.name = name
+        self.type = type

flo_ai-0.0.1/flo_ai/models/flo_planner.py

@@ -0,0 +1,157 @@
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from flo.models.flo_executable import ExecutableFlo
+from flo.state.flo_state import TeamFloAgentStateWithPlan
+from langchain_core.pydantic_v1 import BaseModel, Field
+from typing import Union
+from langgraph.graph.graph import CompiledGraph
+from flo.models.flo_team import FloTeam
+from langgraph.graph import StateGraph
+from flo.state.flo_session import FloSession
+from flo.helpers.utils import randomize_name
+
+
+class Response(BaseModel):
+    """Response to user."""
+    response: str
+
+class Plan(BaseModel):
+    """Plan to follow in future"""
+
+    steps: list[str] = Field(
+        description="different steps to follow, should be in sorted order"
+    )
+
+class Act(BaseModel):
+    """Action to perform."""
+    action: Union[Response, Plan] = Field(
+        description="Action to perform. If you want to respond to user, use Response. "
+        "If you need to further use tools to get the answer, use Plan."
+    )
+
+class FloPlanner(ExecutableFlo):
+    def __init__(self, name: str, graph: CompiledGraph) -> None:
+        super().__init__(name, graph)
+
+    def draw(self, xray=True):
+        return self.graph.get_graph(xray=xray).draw_mermaid_png()
+    
+    def stream(self, work, config = None):
+        return self.graph.stream({
+             "input": work
+        }, config)
+    
+    def invoke(self, work):
+        return self.graph.invoke({
+             "input": work
+        })
+
+class FloPlannerBuilder:
+
+    def __init__(self, 
+                 session: FloSession,
+                 name: str,
+                 team: FloTeam,
+                 planner_prompt: Union[ChatPromptTemplate, None] = None,
+                 replanner_prompt: Union[ChatPromptTemplate, None] = None,
+                 llm: Union[BaseLanguageModel, None] = None) -> None:
+        self.planner_name = randomize_name(name)
+        self.re_planner_name = randomize_name("{}-{}".format("replan", name))
+        self.session = session
+        self.llm = llm if llm is not None else session.llm
+        self.team = team
+        planner_prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    """For the given objective, come up with a simple step by step plan. \
+        This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
+        The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.""",
+                ),
+                MessagesPlaceholder(variable_name="messages"),
+            ]
+        ) if planner_prompt is None else planner_prompt
+        self.planner =  planner_prompt | self.llm.with_structured_output(Plan)
+
+        self.replanner_prompt = ChatPromptTemplate.from_template(
+                """For the given objective, come up with a simple step by step plan. \
+            This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
+            The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.
+
+            Your objective was this:
+            {input}
+
+            Your original plan was this:
+            {plan}
+
+            You have currently done the follow steps:
+            {past_steps}
+
+            Update your plan accordingly. If no more steps are needed and you can return to the user, then do that. 
+            Otherwise if there are still steps to complete, fill out the plan. Only add steps to the plan that still NEED to be done. 
+            Do not return previously done steps as part of the plan."""
+        ) if replanner_prompt is None else replanner_prompt
+        self.replanner =  self.replanner_prompt | self.llm.with_structured_output(Act)
+
+    def execute_step(self, state: TeamFloAgentStateWithPlan):
+        plan = state["plan"]
+        plan_str = "\n".join(f"{i+1}. {step}" for i, step in enumerate(plan))
+        task = plan[0]
+        task_formatted = f"""For the following plan:
+            {plan_str}\n\nYou are tasked with executing step {1}, {task}."""
+        agent_response = self.team.invoke(task_formatted)
+        return {
+            "past_steps": (task, agent_response["messages"][-1].content),
+        }
+
+    def plan_step(self, state: TeamFloAgentStateWithPlan):
+        plan =  self.planner.invoke({"messages": [("user", state["input"])]})
+        if self.session is not None:
+            self.session.append(self.planner_name)
+        return {"plan": plan.steps}
+    
+    def replan_step(self, state: TeamFloAgentStateWithPlan):
+        output = self.replanner.invoke(state)
+        if self.session is not None:
+            self.session.append(self.re_planner_name)
+        if isinstance(output.action, Response):
+            return { "response": output.action.response }
+        else:
+            return { "plan": output.action.steps }
+        
+    def should_end(self, state: TeamFloAgentStateWithPlan):
+        if "response" in state and state["response"]:
+            return "__end__"
+        else:
+            return self.team.name
+
+    def build(self):
+        workflow = StateGraph(TeamFloAgentStateWithPlan)
+
+        # Add the plan node
+        workflow.add_node(self.planner_name, self.plan_step)
+
+        # Add the execution step
+        workflow.add_node(self.team.name, self.execute_step)
+
+        # Add a replan node
+        workflow.add_node(self.re_planner_name, self.replan_step)
+
+        workflow.set_entry_point(self.planner_name)
+
+        # From plan we go to agent
+        workflow.add_edge(self.planner_name, self.team.name)
+
+        # From agent, we replan
+        workflow.add_edge(self.team.name, self.re_planner_name)
+
+        workflow.add_conditional_edges(
+            self.re_planner_name,
+            # Next, we pass in the function that will determine which node is called next.
+            self.should_end,
+        )
+
+        # Finally, we compile it!
+        # This compiles it into a LangChain Runnable,
+        # meaning you can use it as you would any other runnable
+        return FloPlanner(self.planner_name, workflow.compile())

flo_ai-0.0.1/flo_ai/models/flo_rag.py

@@ -0,0 +1,80 @@
+from typing import Optional
+from langchain_core.language_models import BaseLanguageModel
+from langchain.tools import Tool
+from langgraph.graph import END, StateGraph
+from langgraph.graph.graph import CompiledGraph
+from langgraph.prebuilt import ToolNode
+from flo.state.flo_state import TeamFloAgentState
+from langgraph.prebuilt import tools_condition
+from flo.models.flo_executable import ExecutableFlo
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain import hub
+from flo.helpers.utils import randomize_name
+
+class FloRag(ExecutableFlo):
+    def __init__(self, 
+                 name: str, 
+                 graph: CompiledGraph) -> None:
+        super().__init__(name, "team")
+        self.graph = graph
+
+    def draw(self, xray=True):
+        return self.graph.get_graph(xray=xray).draw_mermaid_png()
+
+class FloRagBuilder:
+
+    def __init__(self, 
+                 name: str, 
+                 tools: list[Tool], 
+                 llm: BaseLanguageModel,
+                 prompt: Optional[ChatPromptTemplate] = None) -> None:
+        self.name = randomize_name(name)
+        self.llm = llm
+        self.tools = tools
+        self.prompt = hub.pull("rlm/rag-prompt") if prompt is None else prompt
+    
+    def retriever_agent(self, state: TeamFloAgentState):
+        messages = state["messages"]
+        model = self.llm.bind_tools(self.tools)
+        response = model.invoke(messages)
+        # We return a list, because this will get added to the existing list
+        return { "messages": [response] }
+    
+    def generate(self, state: TeamFloAgentState):
+        messages = state["messages"]
+        question = messages[0].content
+        last_message = messages[-1]
+
+        question = messages[0].content
+        docs = last_message.content
+
+        # Chain
+        rag_chain = self.prompt | self.llm
+
+        # Run
+        response = rag_chain.invoke({"context": docs, "question": question})
+
+        return {"messages": [response] }
+    
+    def build(self) -> FloRag:
+        retrieve = ToolNode(self.tools)
+
+        workflow = StateGraph(TeamFloAgentState)
+        workflow.add_node("agent", self.retriever_agent)
+        workflow.add_node("retrieve", retrieve)  # retrieval
+        workflow.add_node("generate", self.generate)
+        workflow.add_conditional_edges(
+            "agent",
+            tools_condition,
+            {
+                "tools": "retrieve",
+                END: END,
+            },
+        )
+
+        workflow.add_edge("retrieve", "generate")
+        workflow.add_edge("generate", END)
+
+        workflow.set_entry_point("agent")
+        graph = workflow.compile()
+        return FloRag(self.name, graph=graph)

flo_ai-0.0.1/flo_ai/models/flo_supervisor.py

@@ -0,0 +1,97 @@
+from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain.chains import LLMChain
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from typing import Union
+from flo.models.flo_member import FloMember
+from flo.state.flo_session import FloSession
+from flo.constants.prompt_constants import FLO_FINISH
+from flo.helpers.utils import randomize_name
+
+
+# TODO, maybe add description about what team members can do
+supervisor_system_message = (
+    "You are a supervisor tasked with managing a conversation between the"
+    " following {member_type}: {members}. Given the following user request,"
+    " respond with the worker to act next. Each worker will perform a"
+    " task and respond with their results and status. When the users question is answered or the assigned task is finished,"
+    " respond with FINISH. If both your workers has responded atleast once, then strictly return FINISH "
+)
+
+class StateUpdateComponent:
+        def __init__(self, name: str, session: FloSession) -> None:
+            self.name = name
+            self.inner_session = session
+
+        def __call__(self, input):
+            self.inner_session.append(self.name)
+            return input
+
+
+class FloSupervisor:
+    def __init__(self,
+                 executor: LLMChain, 
+                 agents: list[FloMember],
+                 name: str) -> None:
+        self.executor = executor
+        self.agents = agents
+        self.name = name
+        self.type = agents[0].type
+
+    def is_agent_supervisor(self):
+        return self.type == "agent"
+
+class FloSupervisorBuilder:
+    def __init__(self,
+                 session: FloSession,
+                 name: str,
+                 reportees: list[FloMember],
+                 supervisor_prompt: Union[ChatPromptTemplate, None] = None,
+                 llm: Union[BaseLanguageModel, None] = None) -> None:
+        # TODO add validation for reporteess
+        self.name = randomize_name(name)
+        self.session = session
+        self.llm = llm if llm is not None else session.llm
+        self.agents = reportees
+        self.members = [agent.name for agent in reportees]
+        self.options = self.members + [FLO_FINISH]
+        member_type = "workers" if reportees[0].type == "agent" else "team members"
+        self.supervisor_prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", supervisor_system_message),
+                MessagesPlaceholder(variable_name="messages"),
+                (
+                    "system",
+                    "Given the conversation above, who should act next?"
+                    " Or should we FINISH if the task is already answered, (do not wait for the entire question to be answered) or there isn't enough information to answer the question, even then return FINISH? Select one of: {options}",
+                ),
+            ]
+        ).partial(options=str(self.options), members=", ".join(self.members), member_type=member_type) if supervisor_prompt is None else  supervisor_prompt
+    
+    def build(self):
+        function_def = {
+            "name": "route",
+            "description": "Select the next role.",
+            "parameters": {
+                "title": "routeSchema",
+                "type": "object",
+                "properties": {
+                    "next": {
+                        "title": "Next",
+                        "anyOf": [
+                            {"enum": self.options},
+                        ],
+                    }
+                },
+                "required": ["next"],
+            }
+        }
+            
+        chain = (
+            self.supervisor_prompt
+            | self.llm.bind_functions(functions=[function_def], function_call="route")
+            | JsonOutputFunctionsParser()
+            | StateUpdateComponent(self.name, self.session)
+        )
+
+        return FloSupervisor(chain, self.agents, self.name)

flo_ai-0.0.1/flo_ai/models/flo_team.py

@@ -0,0 +1,126 @@
+import functools
+from langchain_core.messages import HumanMessage
+from langgraph.graph.graph import CompiledGraph
+from langgraph.graph import StateGraph, END
+from langchain.agents import AgentExecutor
+
+from flo.models.flo_executable import ExecutableFlo
+from flo.models.flo_agent import FloAgent
+from flo.models.flo_supervisor import FloSupervisor
+from flo.constants.prompt_constants import FLO_FINISH
+from flo.state.flo_state import TeamFloAgentState
+from flo.state.flo_session import FloSession
+from flo.helpers.utils import randomize_name
+
+class FloAgentNode:
+    def __init__(self, 
+                 agent_node: functools.partial, 
+                 name: str) -> None:
+        self.agent_node = agent_node
+        self.name = name
+
+def teamflo_agent_node(state: TeamFloAgentState, agent: AgentExecutor, name: str):
+    result = agent.invoke(state)
+    return {"messages": [HumanMessage(content=result["output"], name=name)]}
+
+# The following functions interoperate between the top level graph state
+# and the state of the research sub-graph
+# this makes it so that the states of each graph don't get intermixed
+def enter_chain(message: str, members: list[str]):
+    results = {
+        "messages": [HumanMessage(content=message)],
+        "team_members": ", ".join(members),
+    }
+    return results
+
+class FloTeam(ExecutableFlo):
+    def __init__(self, name: str, graph: CompiledGraph) -> None:
+        super().__init__(name, graph)
+
+    def draw(self, xray=True):
+        return self.graph.get_graph(xray=xray).draw_mermaid_png()
+
+class FloTeamChain():
+    def __init__(self, name: str, chain) -> None:
+        self.chain = chain
+        self.name = name
+
+class FloTeamBuilder:
+    def __init__(self, 
+                 session: FloSession,
+                 name: str,
+                 supervisor: FloSupervisor) -> None:
+        self.name = randomize_name(name)
+        self.session = session
+        self.flo_agents = supervisor.agents
+        self.flo_supervisor: FloSupervisor = supervisor
+        self.members = list(map(lambda x: x.name, self.flo_agents))
+
+    def build_node(self, flo_agent: FloAgent):
+        agent_func = functools.partial(teamflo_agent_node, agent=flo_agent.executor, name=flo_agent.name)
+        return FloAgentNode(agent_func, flo_agent.name)
+    
+
+    def build(self):
+        if self.flo_supervisor.is_agent_supervisor():
+            return self.build_agent_supervisor_graph()
+        else:
+            return self.build_team_supervisor_graph()
+        
+    def router(self, state: TeamFloAgentState):
+        next = state["next"]
+        conditional_map = {k: k for k in self.members}
+        conditional_map[FLO_FINISH] = END
+        self.session.append(node=next)
+        if self.session.is_looping(node=next):
+            return conditional_map[FLO_FINISH]
+        return conditional_map[next]
+
+
+    def build_agent_supervisor_graph(self) -> FloTeam:
+        flo_agent_nodes = [self.build_node(flo_agent) for flo_agent in self.flo_agents]
+        workflow = StateGraph(TeamFloAgentState)
+        for flo_agent_node in flo_agent_nodes:
+            workflow.add_node(flo_agent_node.name, flo_agent_node.agent_node)
+        workflow.add_node(self.flo_supervisor.name, self.flo_supervisor.executor)
+        for member in self.members:
+            workflow.add_edge(member, self.flo_supervisor.name)
+
+        workflow.add_conditional_edges(self.flo_supervisor.name, self.router)
+
+        workflow.set_entry_point(self.flo_supervisor.name)
+        workflow_graph = workflow.compile()
+        return FloTeam(self.name, workflow_graph)
+    
+    def build_chain_for_teams(self, flo_team: FloTeam):
+        # TODO lets see if we can convert to members
+        return FloTeamChain(flo_team.name, (
+            functools.partial(enter_chain, members=flo_team.graph.nodes)
+            | flo_team.graph
+        ))
+    
+    def get_last_message(self, state: TeamFloAgentState, second = None) -> str:
+        return state["messages"][-1].content
+
+
+    def join_graph(self, response: dict):
+        return {"messages": [response["messages"][-1]]}
+    
+    def build_team_supervisor_graph(self) -> FloTeam:
+        flo_team_entry_chains = [self.build_chain_for_teams(flo_agent) for flo_agent in self.flo_agents]
+        # Define the graph.
+        super_graph = StateGraph(TeamFloAgentState)
+        # First add the nodes, which will do the work
+        for flo_team_chain in flo_team_entry_chains:
+            super_graph.add_node(flo_team_chain.name, self.get_last_message | flo_team_chain.chain | self.join_graph)
+        super_graph.add_node(self.flo_supervisor.name, self.flo_supervisor.executor)
+
+        for member in self.members:
+            super_graph.add_edge(member, self.flo_supervisor.name)
+
+        super_graph.add_conditional_edges(self.flo_supervisor.name, self.router)
+
+        super_graph.set_entry_point(self.flo_supervisor.name)
+        super_graph = super_graph.compile()
+        return FloTeam(self.name, super_graph)
+

flo_ai-0.0.1/flo_ai/retrievers/chromadb_retriever.py

@@ -0,0 +1,31 @@
+from langchain_community.vectorstores import Chroma
+from langchain_core.embeddings import Embeddings
+from flo.retrievers.flo_retriever import FloRetriever
+from typing import Dict, Optional
+
+import chromadb
+from chromadb.config import Settings, DEFAULT_TENANT, DEFAULT_DATABASE
+
+class ChromDBRetriever(FloRetriever):
+    def __init__(self, 
+                 collection: str,
+                 embedding: Embeddings,
+                 type: str = "persistent",
+                 connection_args: Optional[Dict] = None) -> None:
+        if type == "http":
+            client = chromadb.HttpClient(**connection_args, 
+                                         settings=Settings(), 
+                                         tenant=DEFAULT_TENANT, 
+                                         database=DEFAULT_DATABASE)
+        else:
+            client = chromadb.PersistentClient(
+                **connection_args, 
+                settings=Settings(), 
+                tenant=DEFAULT_TENANT, 
+                database=DEFAULT_DATABASE)
+        chromdb = Chroma(
+            client=client,
+            collection_name=collection,
+            embedding_function=embedding
+        )
+        super().__init__(chromdb)

flo_ai-0.0.1/flo_ai/retrievers/flo_retriever.py

@@ -0,0 +1,6 @@
+from langchain_core.vectorstores import VectorStore
+
+class FloRetriever():
+    def __init__(self, vector_store: VectorStore) -> None:
+        self.vector_store = vector_store
+        self.retriver = self.vector_store.as_retriever()

flo_ai-0.0.1/flo_ai/retrievers/mongodb_retriver.py

@@ -0,0 +1,25 @@
+from langchain_mongodb import MongoDBAtlasVectorSearch
+from langchain_core.embeddings import Embeddings
+from pymongo import MongoClient
+
+from flo.retrievers.flo_retriever import FloRetriever
+
+class MongoDBRetriever(FloRetriever):
+    def __init__(self, 
+                 dburl: str,
+                 db_name: str,
+                 collection: str,
+                 embedding_key: str, 
+                 index_name: str,
+                 embedding: Embeddings) -> None:
+        connection_timeout = 60000
+        mongo_client = MongoClient(dburl, connectTimeoutMS=connection_timeout, socketTimeoutMS=connection_timeout)
+        mongo_embedding_collection = (mongo_client
+                                .get_database(db_name)
+                                .get_collection(collection))
+        super().__init__(MongoDBAtlasVectorSearch(
+            collection=mongo_embedding_collection,
+            embedding_key=embedding_key,
+            embedding=embedding,
+            index_name=index_name,
+        ))

flo_ai-0.0.1/flo_ai/state/flo_session.py

@@ -0,0 +1,38 @@
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.tools import BaseTool
+
+class FloSession:
+    def __init__(self, llm: BaseLanguageModel, loop_size: int = 2, max_loop: int = 3) -> None:
+        self.llm = llm
+        self.tools = dict()
+        self.counter = dict()
+        self.navigation: list[str] = list()
+        self.pattern_series = dict()
+        self.loop_size: int = loop_size
+        self.max_loop: int = max_loop
+
+    def register_tool(self, name: str, tool: BaseTool):
+        self.tools[name] = tool
+        return self
+
+    def append(self, node: str) -> int:
+        self.counter[node] = self.counter.get(node, 0) + 1
+        if node in self.navigation:
+            last_known_index = len(self.navigation) - 1 - self.navigation[::-1].index(node)
+            pattern_array = self.navigation[last_known_index: len(self.navigation)]
+            if len(pattern_array) + 1 >= self.loop_size:
+                pattern = "|".join(pattern_array) + "|" + node
+                if node in self.pattern_series:
+                    self.pattern_series[node].append(pattern)
+                else:
+                    self.pattern_series[node] = [pattern]
+        self.navigation.append(node)
+
+    def is_looping(self, node) -> bool:
+        patterns = self.pattern_series[node] if node in self.pattern_series else []
+        if len(patterns) < self.max_loop:
+            return False
+        return patterns[-(self.max_loop):] == [patterns[-1]] * self.max_loop
+
+    def stringify(self):
+        return str(self.counter)

flo_ai-0.0.1/flo_ai/state/flo_state.py

@@ -0,0 +1,19 @@
+from typing import Annotated, List, Sequence, TypedDict
+from langchain_core.messages import BaseMessage
+from typing import List, Tuple, Annotated, TypedDict, Dict
+
+import operator
+
+# The agent state is the input to each node in the graph
+class TeamFloAgentState(TypedDict):
+    # The annotation tells the graph that new messages will always
+    # be added to the current states
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+    # The 'next' field indicates where to route to next
+    next: str
+
+class TeamFloAgentStateWithPlan(TypedDict):
+    input: str
+    plan: List[str]
+    past_steps: Annotated[List[Tuple], operator.add]
+    response: str

flo_ai-0.0.1/flo_ai/tools/big_query_tool.py

@@ -0,0 +1,43 @@
+from typing import Type, Optional
+from google.cloud import bigquery
+from langchain.tools import BaseTool
+from langchain_core.tools import CallbackManagerForToolRun, AsyncCallbackManagerForToolRun
+from langchain.pydantic_v1 import Field, BaseModel
+
+
+class BigQueryToolInput(BaseModel):
+    query: str = Field(description="""
+                                  The sql query to be run on the BigQuery dataset as string and nothing else
+                                  """)
+    
+    
+class BigQueryExecutorTool(BaseTool):
+    name = "BigQueryReaderTool"
+    description = "A tool that can run SQL queries on BigQuery Tables, just pass query to be run as input and nothing else"
+    args_schema: Type[BaseModel] = BigQueryToolInput
+
+    def _run(
+        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
+    ) -> str:
+        """Use the tool."""
+        return self.bigquery_executor(query)
+
+    async def _arun(
+        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
+    ) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("custom_search does not support async")
+    
+    def bigquery_executor(self, sql_query_string: str) -> str:
+            """Executes a BigQuery SQL query and returns formatted results."""
+            try:
+                query = sql_query_string.replace("```sql", "").replace("```", "")
+                client = bigquery.Client()
+                query_job = client.query(query)
+                results = query_job.result()
+                result_str = "Query Results:\n"
+                for row in results:
+                    result_str += f" - {row}\n"
+                return result_str
+            except Exception as e:
+                return f"Failed to run query: {str(e)}"

flo_ai-0.0.1/flo_ai/tools/retriever_tool.py

@@ -0,0 +1,14 @@
+from langchain.tools.retriever import create_retriever_tool
+from langchain_core.retrievers import BaseRetriever
+
+MONGODB_RETRIVER_TOOL = "MongoDBRetreiverTool"
+CHROMADB_RETRIVER_TOOL = "ChromaDBRetreiverTool"
+
+def fetch_retriver_tool_by_name(tool_name: str, 
+                                tool_description: str,
+                                retriever: BaseRetriever):
+    return create_retriever_tool(
+        retriever,
+        tool_name,
+        tool_description
+    )

flo_ai-0.0.1/flo_ai/yaml/flo_team_builder.py

@@ -0,0 +1,86 @@
+from pydantic import BaseModel
+from typing import List
+import yaml
+import re
+from typing import Optional
+from flo.models.exception import FloValidationException
+
+
+KIND_SUPERVISED_TEAM = "FloRoutedTeam"
+KIND_FLO_AGENT = "FloAgent"
+
+yaml_kinds = [
+  KIND_SUPERVISED_TEAM,
+  KIND_FLO_AGENT
+]
+
+class KeyValueArgs(BaseModel):
+    name: str
+    value: str
+
+class FilterArgs(BaseModel):
+    name: str
+    description: str
+    type: str
+
+class ToolConfig(BaseModel):
+    name: str
+    args: Optional[List[KeyValueArgs]] = None
+    properties: Optional[List[KeyValueArgs]] = None
+    filters: Optional[List[FilterArgs]] = None
+
+class AgentConfig(BaseModel):
+    name: str
+    kind: Optional[str] = None
+    prompt: str
+    tools: List[ToolConfig] = []
+
+class RouterConfig(BaseModel):
+    name: str
+    kind: str
+
+class PlannerConfig(BaseModel):
+    name: str
+
+class TeamConfig(BaseModel):
+    name: str
+    agents: Optional[List[AgentConfig]] = None
+    subteams: Optional[List['TeamConfig']] = None
+    router: Optional[RouterConfig] = None
+    planner: Optional[PlannerConfig] = None
+
+class FloRoutedTeamConfig(BaseModel):
+    apiVersion: str
+    kind: str
+    name: str
+    team: TeamConfig
+
+class FloAgentConfig(BaseModel):
+    apiVersion: str
+    kind: str
+    name: str
+    agent: AgentConfig
+
+def to_supervised_team(yaml_str: str) -> FloRoutedTeamConfig:
+    parsed_data = yaml.safe_load(yaml_str)
+    kind = parsed_data["kind"]
+    if kind == KIND_SUPERVISED_TEAM:
+        flo_supervised_team = FloRoutedTeamConfig(**parsed_data)
+        validate_sup_team_config(flo_supervised_team)
+        return flo_supervised_team
+    elif kind == KIND_FLO_AGENT:
+        flo_agent = FloAgentConfig(**parsed_data)
+        validate_sup_team_config(flo_agent)
+        return flo_agent
+    else:
+        raise FloValidationException("Unknown kind: {}".format(kind))
+
+def validate_sup_team_config(flo: FloRoutedTeamConfig):
+    if flo.kind == KIND_FLO_AGENT:
+        return
+    if flo.name is None or not is_valid_name(flo.name):
+        raise FloValidationException("Invalid agent name while creating the flow, expected: [^[a-z][a-z0-9_-]*$]")
+    
+def is_valid_name(s: str) -> bool:
+    pattern = r'^[a-z][a-z0-9_-]*$'
+    return bool(re.match(pattern, s))

flo_ai-0.0.1/pyproject.toml

@@ -0,0 +1,36 @@
+[tool.poetry]
+name = "flo-ai"
+version = "0.0.1"
+description = "A easy way to create structured AI agents"
+authors = ["vizsatiz <vishnu@rootfor.xyz>"]
+license = "MIT"
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = ">=3.10, <3.13"
+langchain = "0.2.1"
+langgraph = "0.0.55"
+ipython = "8.24.0"
+httpx = "0.27.0"
+langchain-community = "0.2.0"
+langchain-experimental = "0.0.59"
+google-cloud-bigquery = "^3.23.1"
+openai = "1.30.5"
+langchain-openai = "^0.1.8"
+pillow = "^10.3.0"
+langchain-mongodb = "^0.1.5"
+langchain-mistralai = "^0.1.7"
+langchain-chroma = "^0.1.1"
+python-dotenv = "^1.0.1"
+
+
+[tool.poetry.group.dev.dependencies]
+langchainhub = "^0.1.17"
+torch = "2.0.0"
+sentence-transformers = "^3.0.0"
+pypdf = "^4.2.0"
+lark = "^1.1.9"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"