Fix formatting

Author: jlowin

examples/asyncpg-memory.py

@@ -1,9 +1,8 @@
+import asyncio
+
 import controlflow as cf
 from controlflow.memory.async_memory import AsyncMemory
-
 from controlflow.memory.providers.postgres import AsyncPostgresMemory
-import asyncio
-
 
 provider = AsyncPostgresMemory(
     database_url="postgresql+psycopg://postgres:postgres@localhost:5432/database",
@@ -33,7 +32,6 @@ async def remember_pet():
     )
 
 
-
 # Create a flow to recall the user's favorite color
 @cf.flow
 async def recall_pet():
@@ -42,6 +40,7 @@ async def recall_pet():
         agents=[agent],
     )
 
+
 async def main():
     print("First flow:")
     await remember_pet()
@@ -51,5 +50,6 @@ async def main():
     print(result)
     return result
 
+
 if __name__ == "__main__":
     asyncio.run(main())

src/controlflow/defaults.py

@@ -6,8 +6,8 @@
 import controlflow.utilities
 import controlflow.utilities.logging
 from controlflow.llm.models import BaseChatModel
-from controlflow.memory.memory import MemoryProvider, get_memory_provider
 from controlflow.memory.async_memory import AsyncMemoryProvider, get_memory_provider
+from controlflow.memory.memory import MemoryProvider, get_memory_provider
 from controlflow.utilities.general import ControlFlowModel
 
 from .agents import Agent
@@ -40,7 +40,9 @@ class Defaults(ControlFlowModel):
     model: Optional[Any]
     history: History
     agent: Agent
-    memory_provider: Optional[Union[MemoryProvider, str]] | Optional[Union[AsyncMemoryProvider, str]]
+    memory_provider: (
+        Optional[Union[MemoryProvider, str]] | Optional[Union[AsyncMemoryProvider, str]]
+    )
 
     # add more defaults here
     def __repr__(self) -> str:

src/controlflow/memory/__init__.py

@@ -1,2 +1,2 @@
 from .memory import Memory
-from .async_memory import AsyncMemory
+from .async_memory import AsyncMemory

src/controlflow/memory/async_memory.py

@@ -96,7 +96,6 @@ def validate_key(cls, v: str) -> str:
             )
         return sanitized
 
-
     async def _configure_provider(self):
         await self.provider.configure(self.key)
         return self
@@ -133,7 +132,6 @@ def get_tools(self) -> List[Tool]:
 def get_memory_provider(provider: str) -> AsyncMemoryProvider:
     logger.debug(f"Loading memory provider: {provider}")
 
-
     # --- async postgres ---
 
     if provider.startswith("async-postgres"):

src/controlflow/memory/providers/postgres.py

@@ -1,27 +1,27 @@
 import uuid
 from typing import Callable, Dict, Optional
 
+# async pg
+import anyio
 import sqlalchemy
 from pgvector.sqlalchemy import Vector
 from pydantic import Field
 from sqlalchemy import Column, String, select, text
 from sqlalchemy.dialects.postgresql import ARRAY
 from sqlalchemy.exc import ProgrammingError
-from sqlalchemy.orm import Session, declarative_base, sessionmaker
-from sqlalchemy_utils import create_database, database_exists
-
-#async pg
-import anyio
 from sqlalchemy.ext.asyncio import (
-    create_async_engine,
     AsyncEngine,
     AsyncSession,
     async_sessionmaker,
+    create_async_engine,
 )
+from sqlalchemy.orm import Session, declarative_base, sessionmaker
+from sqlalchemy_utils import create_database, database_exists
 
 import controlflow
-from controlflow.memory.memory import MemoryProvider
 from controlflow.memory.async_memory import AsyncMemoryProvider
+from controlflow.memory.memory import MemoryProvider
+
 try:
     # For embeddings, we can use langchain_openai or any other library:
     from langchain_openai import OpenAIEmbeddings
@@ -125,7 +125,7 @@ def configure(self, memory_key: str) -> None:
             pool_timeout=self.pool_timeout,
             pool_recycle=self.pool_recycle,
             pool_pre_ping=self.pool_pre_ping,
-            )
+        )
 
         # 2) If DB doesn't exist, create it!
         if not database_exists(engine.url):
@@ -239,7 +239,6 @@ def search(self, memory_key: str, query: str, n: int = 20) -> Dict[str, str]:
         return {row.id: row.text for row in results}
 
 
-
 class AsyncPostgresMemory(AsyncMemoryProvider):
     """
     An async MemoryProvider storing text + embeddings in PostgreSQL
@@ -249,7 +248,7 @@ class AsyncPostgresMemory(AsyncMemoryProvider):
     database_url: str = Field(
         default="postgresql+asyncpg://user:password@localhost:5432/your_database",
         description="Async Postgres URL with the asyncpg driver, e.g. "
-                    "'postgresql+asyncpg://user:pass@host:5432/dbname'."
+        "'postgresql+asyncpg://user:pass@host:5432/dbname'.",
     )
 
     table_name: str = Field(
@@ -261,34 +260,30 @@ class AsyncPostgresMemory(AsyncMemoryProvider):
 
     embedding_dimension: int = Field(
         default=1536,
-        description="Dimension of the embedding vectors. Must match your model output size."
+        description="Dimension of the embedding vectors. Must match your model output size.",
     )
 
     embedding_fn: Callable = Field(
         default_factory=lambda: OpenAIEmbeddings(model="text-embedding-ada-002"),
-        description="Function that turns a string into a numeric vector."
+        description="Function that turns a string into a numeric vector.",
     )
 
     # -- Pool / Engine settings (SQLAlchemy will do the pooling)
     pool_size: int = Field(
-        5,
-        description="Number of permanent connections in the async pool."
+        5, description="Number of permanent connections in the async pool."
     )
     max_overflow: int = Field(
-        10,
-        description="Number of 'overflow' connections if the pool is full."
+        10, description="Number of 'overflow' connections if the pool is full."
     )
     pool_timeout: int = Field(
-        30,
-        description="Seconds to wait for a connection before raising an error."
+        30, description="Seconds to wait for a connection before raising an error."
     )
     pool_recycle: int = Field(
         1800,
-        description="Recycle connections after N seconds to avoid stale connections."
+        description="Recycle connections after N seconds to avoid stale connections.",
     )
     pool_pre_ping: bool = Field(
-        True,
-        description="Check connection health before using from the pool."
+        True, description="Check connection health before using from the pool."
     )
 
     # We'll store an async engine + session factory:
@@ -299,7 +294,7 @@ class AsyncPostgresMemory(AsyncMemoryProvider):
     _table_class_cache: Dict[str, Base] = {}
 
     _configured: bool = False
-    
+
     async def configure(self, memory_key: str) -> None:
         """
         1) Create an async engine.
@@ -347,12 +342,16 @@ async def configure(self, memory_key: str) -> None:
                     # (2) Actually create it (async):
                     def _sync_create(connection):
                         """Helper function to run table creation in sync context."""
-                        Base.metadata.create_all(connection, tables=[memory_model.__table__])
+                        Base.metadata.create_all(
+                            connection, tables=[memory_model.__table__]
+                        )
 
                     try:
                         await conn.run_sync(_sync_create)
                     except ProgrammingError as e:
-                        raise RuntimeError(f"Failed to create table '{table_name}': {e}")
+                        raise RuntimeError(
+                            f"Failed to create table '{table_name}': {e}"
+                        )
 
             # 4) Now that the DB and table are ready, create a session factory
             self._SessionLocal = async_sessionmaker(
@@ -362,7 +361,6 @@ def _sync_create(connection):
 
             self._configured = True
 
-
     def _get_table(self, memory_key: str) -> Base:
         """
         Return or create the dynamic model class for 'memory_{key}' table.
@@ -456,4 +454,4 @@ async def search(self, memory_key: str, query: str, n: int = 20) -> Dict[str, st
             rows = results.all()
 
         # Convert list of Row objects -> dict
-        return {row.id: row.text for row in rows}
+        return {row.id: row.text for row in rows}

src/controlflow/orchestration/prompt_templates.py

@@ -4,8 +4,8 @@
 
 from controlflow.agents.agent import Agent
 from controlflow.flows import Flow
-from controlflow.memory.memory import Memory
 from controlflow.memory.async_memory import AsyncMemory
+from controlflow.memory.memory import Memory
 from controlflow.tasks.task import Task
 from controlflow.tools.tools import Tool
 from controlflow.utilities.general import ControlFlowModel

src/controlflow/tasks/task.py

@@ -37,8 +37,8 @@
 import controlflow
 from controlflow.agents import Agent
 from controlflow.instructions import get_instructions
-from controlflow.memory.memory import Memory
 from controlflow.memory.async_memory import AsyncMemory
+from controlflow.memory.memory import Memory
 from controlflow.tools import Tool, tool
 from controlflow.tools.input import cli_input
 from controlflow.tools.tools import as_tools