self improving agents

Author: samuelcolvin

human-seeded-evals/app/agent.py

@@ -1,11 +1,18 @@
 from __future__ import annotations as _annotations
 
+import os
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
+from typing import AsyncIterator
 
+from cloudkv import AsyncCloudKV
 from pydantic_ai import Agent, RunContext
+from pydantic_ai.models import Model
 
 from .models import TimeRangeInputs, TimeRangeResponse
+from .self_improving_agent import SelfImprovingAgentModel
+from .self_improving_agent_storage import CloudKVStorage
 
 
 @dataclass
@@ -23,13 +30,24 @@ class TimeRangeDeps:
 )
 
 
+@asynccontextmanager
+async def self_improving_model() -> AsyncIterator[SelfImprovingAgentModel]:
+    cloudkv_read_token, cloudkv_write_token = os.environ['CLOUDKV_TOKEN'].split('.')
+    logfire_read_token = os.environ['LOGFIRE_READ_TOKEN']
+    async with AsyncCloudKV(cloudkv_read_token, cloudkv_write_token) as cloudkv:
+        storage = CloudKVStorage(cloudkv)
+        m = SelfImprovingAgentModel('anthropic:claude-sonnet-4-0', storage, logfire_read_token, 'time_range_agent')
+        yield m
+        await m.wait_for_coach()
+
+
 @time_range_agent.instructions
 def inject_current_time(ctx: RunContext[TimeRangeDeps]) -> str:
     """Add the user's current time and timezone in the format 'Friday, November 22, 2024 11:15:14 PST' to context."""
     return f"The user's current time is {ctx.deps.now:%A, %B %d, %Y %H:%M:%S %Z}."
 
 
-async def infer_time_range(inputs: TimeRangeInputs) -> TimeRangeResponse:
+async def infer_time_range(inputs: TimeRangeInputs, *, model: Model | None = None) -> TimeRangeResponse:
     """Infer a time range from a user prompt."""
-    result = await time_range_agent.run(inputs.prompt, deps=TimeRangeDeps(now=inputs.now))
+    result = await time_range_agent.run(inputs.prompt, deps=TimeRangeDeps(now=inputs.now), model=model)
     return result.output

human-seeded-evals/app/main.py

@@ -1,16 +1,30 @@
+from contextlib import asynccontextmanager
+from typing import cast
+
 import logfire
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 
-from .agent import infer_time_range
+from .agent import infer_time_range, self_improving_model
 from .models import TimeRangeInputs, TimeRangeResponse
+from .self_improving_agent import SelfImprovingAgentModel
 
 logfire.configure(environment='dev')
+
 logfire.instrument_pydantic_ai()
 
-app = FastAPI()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    async with self_improving_model() as model:
+        app.state.model = model
+        yield
+
+
+app = FastAPI(lifespan=lifespan)
 logfire.instrument_fastapi(app)
 
 
 @app.post('/api/timerange')
-async def convert_time_range(time_range_inputs: TimeRangeInputs) -> TimeRangeResponse:
-    return await infer_time_range(time_range_inputs)
+async def convert_time_range(request: Request, time_range_inputs: TimeRangeInputs) -> TimeRangeResponse:
+    model = cast(SelfImprovingAgentModel, request.app.state.model)
+    return await infer_time_range(time_range_inputs, model=model)