Code formatting using Ruff

.github/workflows/build-project.yml

@@ -29,8 +29,12 @@ jobs:
     - name: Install dependencies
       run: poetry lock && poetry install
 
+    - name: Run Ruff
+      run: poetry run pre-commit run --all-files
+
     - name: Run build
       run: poetry build
 
     - name: Run tests
-      run: poetry run pytest
+      run: poetry run pytest
+

.pre-commit-config.yaml

@@ -0,0 +1,12 @@
+repos:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.7.3
+  hooks:
+    # Run the linter.
+    - id: ruff
+      files: \.py$
+    # Run the formatter.
+    - id: ruff-format
+      files: \.py$
+      args: [--config, format.quote-style = 'single']

examples/delegator_example.py

@@ -3,6 +3,7 @@
 from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search.tool import TavilySearchResults
 from dotenv import load_dotenv
+
 load_dotenv()
 
 yaml_data = """
@@ -41,11 +42,10 @@
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
 session = FloSession(llm).register_tool(
-    name="TavilySearchResults", 
-    tool=TavilySearchResults()
+    name='TavilySearchResults', tool=TavilySearchResults()
 )
 
 flo: Flo = Flo.build(session, yaml=yaml_data)
-flo.draw_to_file("delegate.png", xray=True)
+flo.draw_to_file('delegate.png', xray=True)
 # data = flo.invoke(input_prompt)
-# print((data['messages'][-1]).content)
+# print((data['messages'][-1]).content)

examples/hierarchical_blogging_team.py

@@ -1,9 +1,9 @@
-from flo_ai import Flo
 from flo_ai.core import Flo
 from langchain_openai import ChatOpenAI
 from flo_ai import FloSession
 from langchain_community.tools.tavily_search.tool import TavilySearchResults
 from dotenv import load_dotenv
+
 load_dotenv()
 
 yaml_data = """
@@ -47,6 +47,6 @@
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o')
 session = FloSession(llm).register_tool(
-    name="TavilySearchResults", tool=TavilySearchResults()
+    name='TavilySearchResults', tool=TavilySearchResults()
 )
-flo: Flo = Flo.build(session, yaml=yaml_data)
+flo: Flo = Flo.build(session, yaml=yaml_data)

examples/linear_router_team.py

@@ -1,36 +1,44 @@
-from flo_ai import Flo
 from flo_ai.core import Flo
 from langchain_openai import ChatOpenAI
 from flo_ai import FloSession
 from langchain_community.tools.tavily_search.tool import TavilySearchResults
 from dotenv import load_dotenv
-load_dotenv()
+
 
 from typing import Optional, Type
 from pydantic import BaseModel, Field
-from langchain.tools import BaseTool, StructuredTool, tool
+from langchain.tools import BaseTool
 from langchain.callbacks.manager import (
     AsyncCallbackManagerForToolRun,
     CallbackManagerForToolRun,
 )
 
+load_dotenv()
+
+
 class FetchTrxInput(BaseModel):
-    reference_number: str = Field(description="The transaction reference number")
+    reference_number: str = Field(description='The transaction reference number')
+
 
 class FetchTransactionTool(BaseTool):
-    name = "fetch_transactions"
-    description = "useful for when you want to fetch the transaction details given reference number"
+    name = 'fetch_transactions'
+    description = 'useful for when you want to fetch the transaction details given reference number'
     args_schema: Type[BaseModel] = FetchTrxInput
 
     def _run(
-        self, reference_number: str, run_manager: Optional[CallbackManagerForToolRun] = None
+        self,
+        reference_number: str,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
     ) -> str:
-        return "The transaction happened on 23/07/2024 IST and it failed because there was not enough balance in the account"
+        return 'The transaction happened on 23/07/2024 IST and it failed because there was not enough balance in the account'
 
     async def _arun(
-        self, reference_number: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
+        self,
+        reference_number: str,
+        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
     ) -> str:
-        return "The transaction happened on 23/07/2024 IST and it failed because there was not enough balance in the account"
+        return 'The transaction happened on 23/07/2024 IST and it failed because there was not enough balance in the account'
+
 
 yaml_data = """
 apiVersion: flo/alpha-v1
@@ -59,11 +67,11 @@ async def _arun(
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o')
 session = FloSession(llm).register_tool(
-    name="TavilySearchResults", tool=TavilySearchResults()
+    name='TavilySearchResults', tool=TavilySearchResults()
 )
 flo: Flo = Flo.build(session, yaml=yaml_data)
 
 for event in flo.stream(input_prompt):
     for k, v in event.items():
-        if k != "__end__":
+        if k != '__end__':
             print(v)

examples/llm_extensibility.py

@@ -1,34 +1,44 @@
-import os
 from flo_ai import Flo
 from flo_ai import FloSession
 from pydantic import BaseModel, Field
-from langchain_openai import ChatOpenAI, AzureChatOpenAI
+from langchain_openai import ChatOpenAI
 from flo_ai.tools.flo_tool import flotool
 
 from dotenv import load_dotenv
+import warnings
+
 load_dotenv()
 
-import warnings
+
 warnings.simplefilter('default', DeprecationWarning)
 
 gpt35 = ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo')
 gpt_4o_mini = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
 gpt_4o = ChatOpenAI(temperature=0, model_name='gpt-4o')
 session = FloSession(gpt35)
 
-session.register_model("bronze", gpt35)
-session.register_model("silver", gpt_4o_mini)
-session.register_model("gold", gpt_4o)
+session.register_model('bronze', gpt35)
+session.register_model('silver', gpt_4o_mini)
+session.register_model('gold', gpt_4o)
+
 
 class SendEmailInput(BaseModel):
-    to: str = Field(description="Comma seperared list of users emails to which email needs to be sent")
-    message: str = Field(description="The email text to be sent")
+    to: str = Field(
+        description='Comma seperared list of users emails to which email needs to be sent'
+    )
+    message: str = Field(description='The email text to be sent')
+
 
-@flotool("email_triage", "useful for when you need to send an email to someone", argument_contract=SendEmailInput)
+@flotool(
+    'email_triage',
+    'useful for when you need to send an email to someone',
+    argument_contract=SendEmailInput,
+)
 def email_tool(to: str, message: str):
-    return f"Email sent successfully to: {to}"
+    return f'Email sent successfully to: {to}'
 
-session.register_tool("SendEmailTool", email_tool)
+
+session.register_tool('SendEmailTool', email_tool)
 
 agent_yaml = """
 apiVersion: flo/alpha-v1
@@ -70,10 +80,6 @@ def email_tool(to: str, message: str):
 
 flo: Flo = Flo.build(session, yaml=agent_yaml)
 for s in flo.stream(input_prompt):
-     if "__end__" not in s:
+    if '__end__' not in s:
         print(s)
-        print("----")
-
-
-
-
+        print('----')

examples/rag_tool.py

@@ -1,7 +1,6 @@
 from flo_ai import Flo
 from flo_ai import FloSession
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
 from langchain_community.document_loaders import TextLoader
 from langchain_community.embeddings.sentence_transformer import (
@@ -10,46 +9,47 @@
 from langchain_text_splitters import CharacterTextSplitter
 
 from dotenv import load_dotenv
-load_dotenv()
+from flo_ai.retrievers.flo_retriever import FloRagBuilder
+from flo_ai.retrievers.flo_compression_pipeline import FloCompressionPipeline
+
 
+load_dotenv()
 
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
 
-session = FloSession(
-    llm, 
-    log_level="ERROR"
-)
+session = FloSession(llm, log_level='ERROR')
 
 # load the document and split it into chunks
-loader = TextLoader("./examples/rag_document.txt")
+loader = TextLoader('./examples/rag_document.txt')
 documents = loader.load()
 
 # split it into chunks
 text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
 docs = text_splitter.split_documents(documents)
 
 # create the open-source embedding function
-embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+embedding_function = SentenceTransformerEmbeddings(model_name='all-MiniLM-L6-v2')
 
 # load it into Chroma
 db = Chroma.from_documents(docs, embedding_function)
 
-from flo_ai.retrievers.flo_retriever import FloRagBuilder
-from flo_ai.retrievers.flo_compression_pipeline import FloCompressionPipeline
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
 session = FloSession(llm)
 builder = FloRagBuilder(session, db.as_retriever())
-compression_pipeline = FloCompressionPipeline(OpenAIEmbeddings(model="text-embedding-3-small"))
+compression_pipeline = FloCompressionPipeline(
+    OpenAIEmbeddings(model='text-embedding-3-small')
+)
 compression_pipeline.add_embedding_reduntant_filter()
 compression_pipeline.add_embedding_relevant_filter()
 # Reranking
 
-retriever_tool = builder.with_compression(compression_pipeline).build_rag_tool(name="HousingLoanRetreiver",
-                                                                                               description="Tool to fetch data around housing loans")
-session.register_tool(name="HousingLoanTool", tool=retriever_tool)
-
+retriever_tool = builder.with_compression(compression_pipeline).build_rag_tool(
+    name='HousingLoanRetreiver', description='Tool to fetch data around housing loans'
+)
+session.register_tool(name='HousingLoanTool', tool=retriever_tool)
+
 simple_tool_agent = """
 apiVersion: flo/alpha-v1
 kind: FloAgent
@@ -64,4 +64,4 @@
 
 flo = Flo.build(session, simple_tool_agent)
 
-print(flo.invoke("Whats interest rate on loan"))
+print(flo.invoke('Whats interest rate on loan'))

examples/rag_with_reranking.py

@@ -4,58 +4,67 @@
 from langchain_openai import OpenAIEmbeddings
 from langchain_openai import ChatOpenAI
 from dotenv import load_dotenv
-load_dotenv()
+
 
 from flo_ai import FloSession
 from flo_ai.retrievers.flo_retriever import FloRagBuilder
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from flo_ai.retrievers.flo_compression_pipeline import FloCompressionPipeline
+import logging
 
-db_url = os.getenv("MONGO_DB_URL")
+
+load_dotenv()
+db_url = os.getenv('MONGO_DB_URL')
 
 connection_timeout = 60000
-mongo_client = MongoClient(db_url, connectTimeoutMS=connection_timeout, socketTimeoutMS=connection_timeout)
-mongo_embedding_collection = (mongo_client
-                        .get_database("dohabank")
-                        .get_collection("products"))
+mongo_client = MongoClient(
+    db_url, connectTimeoutMS=connection_timeout, socketTimeoutMS=connection_timeout
+)
+mongo_embedding_collection = mongo_client.get_database('dohabank').get_collection(
+    'products'
+)
 
 store = MongoDBAtlasVectorSearch(
     collection=mongo_embedding_collection,
-    embedding_key="embedding",
-    embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
-    index_name="bank-products-index",
+    embedding_key='embedding',
+    embedding=OpenAIEmbeddings(model='text-embedding-3-small'),
+    index_name='bank-products-index',
 )
 
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
 session = FloSession(llm)
 rag_builder = FloRagBuilder(session, store.as_retriever())
 
-import logging
 
 logging.basicConfig()
-logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)
+logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)
 
 custom_prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", """You are an assistant for banking employees, of Doha Bank.
+    [
+        (
+            'system',
+            """You are an assistant for banking employees, of Doha Bank.
                  Use the following pieces of retrieved context to answer the question. 
                  If you don't know the answer, just say that you don't know. 
-                 Try to answer questions as bullet points that are easy to read"""),
-                MessagesPlaceholder(variable_name="chat_history"),
-                ("human", "{question}"),
-            ]
-        )
-
-from langchain.schema import BaseMessage
-compression_pipeline = FloCompressionPipeline(OpenAIEmbeddings(model="text-embedding-3-small"))
+                 Try to answer questions as bullet points that are easy to read""",
+        ),
+        MessagesPlaceholder(variable_name='chat_history'),
+        ('human', '{question}'),
+    ]
+)
+
+
+compression_pipeline = FloCompressionPipeline(
+    OpenAIEmbeddings(model='text-embedding-3-small')
+)
 compression_pipeline.add_embedding_reduntant_filter()
 compression_pipeline.add_embedding_relevant_filter()
 
-rag = rag_builder.with_prompt(
-    custom_prompt
-).with_multi_query().with_compression(
-    compression_pipeline
-    ).build_rag()
-print(rag.invoke({ "question": "What are the documents applying for housing loan" }))
-
+rag = (
+    rag_builder.with_prompt(custom_prompt)
+    .with_multi_query()
+    .with_compression(compression_pipeline)
+    .build_rag()
+)
+print(rag.invoke({'question': 'What are the documents applying for housing loan'}))