Merge pull request #62 from jlowin/litellm

Add liteagent
PrefectHQ · May 21, 2024 · 32dca0b · 32dca0b
2 parents 98d11db + 586dceb
commit 32dca0b
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 80 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -12,7 +12,6 @@ dependencies = [
     "pydantic>=2",
     "textual>=0.61.1",
     "litellm>=1.37.17",
-    "numpydoc>=1.7.0",
 ]
 readme = "README.md"
 requires-python = ">= 3.9"

diff --git a/requirements-dev.lock b/requirements-dev.lock
@@ -14,8 +14,6 @@ aiosignal==1.3.1
     # via aiohttp
 aiosqlite==0.20.0
     # via prefect
-alabaster==0.7.16
-    # via sphinx
 alembic==1.13.1
     # via prefect
 annotated-types==0.6.0
@@ -42,7 +40,6 @@ attrs==23.2.0
     # via referencing
 babel==2.15.0
     # via mkdocs-material
-    # via sphinx
 boto3==1.34.103
     # via moto
 botocore==1.34.103
@@ -118,8 +115,6 @@ distro==1.9.0
     # via openai
 docker==6.1.3
     # via prefect
-docutils==0.21.2
-    # via sphinx
 envier==0.5.1
     # via ddtrace
 execnet==2.1.1
@@ -177,8 +172,6 @@ idna==3.6
     # via httpx
     # via requests
     # via yarl
-imagesize==1.4.1
-    # via sphinx
 importlib-metadata==7.0.0
     # via litellm
     # via mike
@@ -204,7 +197,6 @@ jinja2==3.1.3
     # via mkdocstrings
     # via moto
     # via prefect
-    # via sphinx
 jinja2-humanize-extension==0.4.0
     # via prefect
 jmespath==1.0.1
@@ -292,8 +284,6 @@ nodeenv==1.8.0
     # via pre-commit
 numpy==1.26.4
     # via prefect
-numpydoc==1.7.0
-    # via controlflow
 oauthlib==3.2.2
     # via kubernetes
     # via requests-oauthlib
@@ -310,7 +300,6 @@ packaging==24.0
     # via mkdocs
     # via prefect
     # via pytest
-    # via sphinx
 paginate==0.5.6
     # via mkdocs-material
 parso==0.8.3
@@ -372,7 +361,6 @@ pygments==2.17.2
     # via ipython
     # via mkdocs-material
     # via rich
-    # via sphinx
 pymdown-extensions==10.8.1
     # via mkdocs-material
     # via mkdocstrings
@@ -459,7 +447,6 @@ requests==2.31.0
     # via prefect
     # via requests-oauthlib
     # via responses
-    # via sphinx
     # via tiktoken
 requests-oauthlib==2.0.0
     # via apprise
@@ -508,22 +495,6 @@ sniffio==1.3.1
     # via httpx
     # via openai
     # via prefect
-snowballstemmer==2.2.0
-    # via sphinx
-sphinx==7.3.7
-    # via numpydoc
-sphinxcontrib-applehelp==1.0.8
-    # via sphinx
-sphinxcontrib-devhelp==1.0.6
-    # via sphinx
-sphinxcontrib-htmlhelp==2.0.5
-    # via sphinx
-sphinxcontrib-jsmath==1.0.1
-    # via sphinx
-sphinxcontrib-qthelp==1.0.7
-    # via sphinx
-sphinxcontrib-serializinghtml==1.1.10
-    # via sphinx
 sqlalchemy==2.0.29
     # via alembic
     # via prefect
@@ -533,8 +504,6 @@ stack-data==0.6.3
     # via ipython
 starlette==0.36.3
     # via fastapi
-tabulate==0.9.0
-    # via numpydoc
 text-unidecode==1.3
     # via python-slugify
 textual==0.61.1

diff --git a/requirements.lock b/requirements.lock
@@ -14,8 +14,6 @@ aiosignal==1.3.1
     # via aiohttp
 aiosqlite==0.20.0
     # via prefect
-alabaster==0.7.16
-    # via sphinx
 alembic==1.13.1
     # via prefect
 annotated-types==0.6.0
@@ -42,7 +40,6 @@ attrs==23.2.0
     # via referencing
 babel==2.15.0
     # via mkdocs-material
-    # via sphinx
 boto3==1.34.103
     # via moto
 botocore==1.34.103
@@ -118,8 +115,6 @@ distro==1.9.0
     # via openai
 docker==6.1.3
     # via prefect
-docutils==0.21.2
-    # via sphinx
 envier==0.5.1
     # via ddtrace
 execnet==2.1.1
@@ -177,8 +172,6 @@ idna==3.6
     # via httpx
     # via requests
     # via yarl
-imagesize==1.4.1
-    # via sphinx
 importlib-metadata==7.0.0
     # via litellm
     # via mike
@@ -204,7 +197,6 @@ jinja2==3.1.3
     # via mkdocstrings
     # via moto
     # via prefect
-    # via sphinx
 jinja2-humanize-extension==0.4.0
     # via prefect
 jmespath==1.0.1
@@ -292,8 +284,6 @@ nodeenv==1.8.0
     # via pre-commit
 numpy==1.26.4
     # via prefect
-numpydoc==1.7.0
-    # via controlflow
 oauthlib==3.2.2
     # via kubernetes
     # via requests-oauthlib
@@ -310,7 +300,6 @@ packaging==24.0
     # via mkdocs
     # via prefect
     # via pytest
-    # via sphinx
 paginate==0.5.6
     # via mkdocs-material
 parso==0.8.4
@@ -372,7 +361,6 @@ pygments==2.17.2
     # via ipython
     # via mkdocs-material
     # via rich
-    # via sphinx
 pymdown-extensions==10.8.1
     # via mkdocs-material
     # via mkdocstrings
@@ -459,7 +447,6 @@ requests==2.31.0
     # via prefect
     # via requests-oauthlib
     # via responses
-    # via sphinx
     # via tiktoken
 requests-oauthlib==2.0.0
     # via apprise
@@ -508,22 +495,6 @@ sniffio==1.3.1
     # via httpx
     # via openai
     # via prefect
-snowballstemmer==2.2.0
-    # via sphinx
-sphinx==7.3.7
-    # via numpydoc
-sphinxcontrib-applehelp==1.0.8
-    # via sphinx
-sphinxcontrib-devhelp==1.0.6
-    # via sphinx
-sphinxcontrib-htmlhelp==2.0.5
-    # via sphinx
-sphinxcontrib-jsmath==1.0.1
-    # via sphinx
-sphinxcontrib-qthelp==1.0.7
-    # via sphinx
-sphinxcontrib-serializinghtml==1.1.10
-    # via sphinx
 sqlalchemy==2.0.29
     # via alembic
     # via prefect
@@ -533,8 +504,6 @@ stack-data==0.6.3
     # via ipython
 starlette==0.36.3
     # via fastapi
-tabulate==0.9.0
-    # via numpydoc
 text-unidecode==1.3
     # via python-slugify
 textual==0.61.1

diff --git a/src/controlflow/core/agent.py b/src/controlflow/core/agent.py
@@ -1,11 +1,13 @@
 import logging
-from typing import Union
+from typing import Callable, Optional, Union
 
+from litellm import Message
 from marvin.utilities.asyncio import ExposeSyncMethodsMixin, expose_sync_method
 from pydantic import Field
 
 from controlflow.core.flow import Flow, get_flow
 from controlflow.core.task import Task
+from controlflow.llm.completions import Response, completion, completion_async
 from controlflow.tools.talk_to_human import talk_to_human
 from controlflow.utilities.prefect import (
     wrap_prefect_tool,
@@ -62,3 +64,51 @@ async def run_async(
 
     def __hash__(self):
         return id(self)
+
+
+class LiteAgent(ControlFlowModel, ExposeSyncMethodsMixin):
+    name: str = Field(
+        ...,
+        description="The name of the agent. This is used to identify the agent in the system and should be unique per assigned task.",
+    )
+    description: Optional[str] = Field(
+        None, description="A description of the agent, visible to other agents."
+    )
+    instructions: Optional[str] = Field(
+        None, description="Instructions for the agent, private to this agent."
+    )
+    tools: list[Callable] = Field(
+        [], description="List of tools availble to the agent."
+    )
+    user_access: bool = Field(
+        False,
+        description="If True, the agent is given tools for interacting with a human user.",
+    )
+    model: Optional[str] = Field(
+        None,
+        description="The model used by the agent. If not provided, the default model will be used.",
+    )
+
+    async def say_async(self, messages: Union[str, dict]) -> Response:
+        if not isinstance(messages, list):
+            raise ValueError("Messages must be provided as a list.")
+
+        messages = [
+            Message(role="user", content=m) if isinstance(m, str) else m
+            for m in messages
+        ]
+
+        return await completion_async(
+            messages=messages, model=self.model, tools=self.tools
+        )
+
+    async def say(self, messages: Union[str, dict]) -> Response:
+        if not isinstance(messages, list):
+            raise ValueError("Messages must be provided as a list.")
+
+        messages = [
+            Message(role="user", content=m) if isinstance(m, str) else m
+            for m in messages
+        ]
+
+        return completion(messages=messages, model=self.model, tools=self.tools)
diff --git a/src/controlflow/llm/completions.py b/src/controlflow/llm/completions.py
@@ -39,15 +39,15 @@ def completion(
     Returns:
         A litellm.ModelResponse object representing the completion response.
     """
+
     intermediate_messages = []
     intermediate_responses = []
 
     if model is None:
         model = controlflow.settings.model
-    if tools is not None:
-        tool_dicts = [function_to_tool_dict(tool) for tool in tools]
-    else:
-        tool_dicts = None
+
+    tool_dicts = [function_to_tool_dict(tool) for tool in tools or []] or None
+
     response = litellm.completion(
         model=model,
         messages=messages,
@@ -101,10 +101,7 @@ def stream_completion(
     if model is None:
         model = controlflow.settings.model
 
-    if tools is not None:
-        tool_dicts = [function_to_tool_dict(tool) for tool in tools]
-    else:
-        tool_dicts = None
+    tool_dicts = [function_to_tool_dict(tool) for tool in tools or []] or None
 
     chunks = []
     for chunk in litellm.completion(
@@ -163,10 +160,7 @@ async def completion_async(
     if model is None:
         model = controlflow.settings.model
 
-    if tools is not None:
-        tool_dicts = [function_to_tool_dict(tool) for tool in tools]
-    else:
-        tool_dicts = None
+    tool_dicts = [function_to_tool_dict(tool) for tool in tools or []] or None
 
     response = await litellm.acompletion(
         model=model,
@@ -221,10 +215,7 @@ async def stream_completion_async(
     if model is None:
         model = controlflow.settings.model
 
-    if tools is not None:
-        tool_dicts = [function_to_tool_dict(tool) for tool in tools]
-    else:
-        tool_dicts = None
+    tool_dicts = [function_to_tool_dict(tool) for tool in tools or []] or None
 
     chunks = []
     async for chunk in litellm.acompletion(