agent: mypy type fixes

mypy examples/agent/__main__.py
mypy examples/agent/fastify.py
mypy examples/openai/__main__.py

Author: Olivier Chafik

examples/agent/agent.py

@@ -3,7 +3,7 @@
 from time import sleep
 import typer
 from pydantic import BaseModel, Json, TypeAdapter
-from typing import Annotated, Callable, List, Union, Optional, Type
+from typing import Annotated, Any, Callable, Dict, List, Union, Optional, Type
 import json, requests
 
 from examples.agent.openapi_client import OpenAPIMethod, openapi_methods_from_endpoint
@@ -13,7 +13,7 @@
 from examples.openai.prompting import ToolsPromptStyle
 from examples.openai.subprocesses import spawn_subprocess
 
-def _get_params_schema(fn: Callable, verbose):
+def _get_params_schema(fn: Callable[[Any], Any], verbose):
     if isinstance(fn, OpenAPIMethod):
         return fn.parameters_schema
 
@@ -26,9 +26,9 @@ def _get_params_schema(fn: Callable, verbose):
 
 def completion_with_tool_usage(
         *,
-        response_model: Optional[Union[Json, Type]]=None,
+        response_model: Optional[Union[Json[Any], type]]=None,
         max_iterations: Optional[int]=None,
-        tools: List[Callable],
+        tools: List[Callable[..., Any]],
         endpoint: str,
         messages: List[Message],
         auth: Optional[str],
@@ -56,7 +56,7 @@ def completion_with_tool_usage(
             type="function",
             function=ToolFunction(
                 name=fn.__name__,
-                description=fn.__doc__,
+                description=fn.__doc__ or '',
                 parameters=_get_params_schema(fn, verbose=verbose)
             )
         )
@@ -128,15 +128,15 @@ def completion_with_tool_usage(
 def main(
     goal: Annotated[str, typer.Option()],
     tools: Optional[List[str]] = None,
-    format: Annotated[str, typer.Option(help="The output format: either a Python type (e.g. 'float' or a Pydantic model defined in one of the tool files), or a JSON schema, e.g. '{\"format\": \"date\"}'")] = None,
+    format: Annotated[Optional[str], typer.Option(help="The output format: either a Python type (e.g. 'float' or a Pydantic model defined in one of the tool files), or a JSON schema, e.g. '{\"format\": \"date\"}'")] = None,
     max_iterations: Optional[int] = 10,
     std_tools: Optional[bool] = False,
     auth: Optional[str] = None,
     parallel_calls: Optional[bool] = False,
     verbose: bool = False,
     style: Optional[ToolsPromptStyle] = None,
 
-    model: Annotated[Optional[Path], typer.Option("--model", "-m")] = "models/7B/ggml-model-f16.gguf",
+    model: Annotated[str, typer.Option("--model", "-m")] = "models/7B/ggml-model-f16.gguf",
     endpoint: Optional[str] = None,
     context_length: Optional[int] = None,
     # endpoint: str = 'http://localhost:8080/v1/chat/completions',
@@ -187,8 +187,8 @@ def main(
         sleep(5)
 
     tool_functions = []
-    types = {}
-    for f in tools:
+    types: Dict[str, type] = {}
+    for f in (tools or []):
         if f.startswith('http://') or f.startswith('https://'):
             tool_functions.extend(openapi_methods_from_endpoint(f))
         else:
@@ -203,7 +203,7 @@ def main(
     if std_tools:
         tool_functions.extend(collect_functions(StandardTools))
 
-    response_model = None #str
+    response_model: Union[type, Json[Any]] = None #str
     if format:
         if format in types:
             response_model = types[format]
@@ -246,10 +246,7 @@ def main(
         seed=seed,
         n_probs=n_probs,
         min_keep=min_keep,
-        messages=[{
-            "role": "user",
-            "content": goal,
-        }]
+        messages=[Message(role="user", content=goal)],
     )
     print(result if response_model else f'➡️ {result}')
     # exit(0)

examples/agent/fastify.py

@@ -17,7 +17,7 @@ def bind_functions(app, module):
         if k == k.capitalize():
             continue
         v = getattr(module, k)
-        if not callable(v) or isinstance(v, Type):
+        if not callable(v) or isinstance(v, type):
             continue
         if not hasattr(v, '__annotations__'):
             continue

examples/agent/openapi_client.py

@@ -17,43 +17,44 @@ def __init__(self, url, name, descriptor, catalog):
         request_body = post_descriptor.get('requestBody')
 
         self.parameters = {p['name']: p for p in parameters}
-        assert all(param['in'] == 'query' for param in self.parameters.values()), f'Only query path parameters are supported (path: {path}, descriptor: {json.dumps(descriptor)})'
+        assert all(param['in'] == 'query' for param in self.parameters.values()), f'Only query path parameters are supported (path: {url}, descriptor: {json.dumps(descriptor)})'
 
         self.body = None
-        self.body_name = None
         if request_body:
-            assert 'application/json' in request_body['content'], f'Only application/json is supported for request body (path: {path}, descriptor: {json.dumps(descriptor)})'
+            assert 'application/json' in request_body['content'], f'Only application/json is supported for request body (path: {url}, descriptor: {json.dumps(descriptor)})'
+
+            body_name = 'body'
+            i = 2
+            while body_name in self.parameters:
+                body_name = f'body{i}'
+                i += 1
+
             self.body = dict(
+                name=body_name,
                 required=request_body['required'],
                 schema=request_body['content']['application/json']['schema'],
             )
 
-            self.body_name = 'body'
-            i = 2
-            while self.body_name in self.parameters:
-                self.body_name = f'body{i}'
-                i += 1
-
         self.parameters_schema = dict(
             type='object',
             properties={
                 **({
-                    self.body_name: self.body['schema']
+                    self.body['name']: self.body['schema']
                 } if self.body else {}),
                 **{
                     name: param['schema']
                     for name, param in self.parameters.items()
                 }
             },
             components=catalog.get('components'),
-            required=[name for name, param in self.parameters.items() if param['required']] + ([self.body_name] if self.body and self.body['required'] else [])
+            required=[name for name, param in self.parameters.items() if param['required']] + ([self.body['name']] if self.body and self.body['required'] else [])
         )
 
     def __call__(self, **kwargs):
         if self.body:
-            body = kwargs.pop(self.body_name, None)
+            body = kwargs.pop(self.body['name'], None)
             if self.body['required']:
-                assert body is not None, f'Missing required body parameter: {self.body_name}'
+                assert body is not None, f'Missing required body parameter: {self.body["name"]}'
         else:
             body = None
 

examples/agent/tools/std_tools.py

@@ -15,6 +15,9 @@ class Duration(BaseModel):
     months: Optional[int] = None
     years: Optional[int] = None
 
+    def __str__(self) -> str:
+        return f"{self.years} years, {self.months} months, {self.days} days, {self.hours} hours, {self.minutes} minutes, {self.seconds} seconds"
+
     @property
     def get_total_seconds(self) -> int:
         return sum([
@@ -29,6 +32,10 @@ def get_total_seconds(self) -> int:
 class WaitForDuration(BaseModel):
     duration: Duration
 
+    def __call__(self):
+        sys.stderr.write(f"Waiting for {self.duration}...\n")
+        time.sleep(self.duration.get_total_seconds)
+
 class WaitForDate(BaseModel):
     until: date
 
@@ -43,7 +50,7 @@ def __call__(self):
 
         days, seconds = time_diff.days, time_diff.seconds
 
-        sys.stderr.write(f"Waiting for {days} days and {seconds} seconds until {d}...\n")
+        sys.stderr.write(f"Waiting for {days} days and {seconds} seconds until {self.until}...\n")
         time.sleep(days * 86400 + seconds)
         sys.stderr.write(f"Reached the target date: {self.until}\n")
 
@@ -67,8 +74,8 @@ def wait(_for: Union[WaitForDuration, WaitForDate]) -> None:
         return _for()
 
     @staticmethod
-    def say_out_loud(something: str) -> str:
+    def say_out_loud(something: str) -> None:
         """
             Just says something. Used to say each thought out loud
         """
-        return subprocess.check_call(["say", something])
+        subprocess.check_call(["say", something])

examples/agent/utils.py

@@ -9,8 +9,10 @@ def load_source_as_module(source):
         i += 1
 
     spec = importlib.util.spec_from_file_location(module_name, source)
+    assert spec, f'Failed to load {source} as module'
     module = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = module
+    assert spec.loader, f'{source} spec has no loader'
     spec.loader.exec_module(module)
     return module
 
@@ -29,7 +31,7 @@ def collect_functions(module):
         if k == k.capitalize():
             continue
         v = getattr(module, k)
-        if not callable(v) or isinstance(v, Type):
+        if not callable(v) or isinstance(v, type):
             continue
         if not hasattr(v, '__annotations__'):
             continue

examples/json_schema_to_grammar.py

@@ -55,9 +55,9 @@ def opt_repetitions(up_to_n, prefix_with_sep=False):
 
 
 class BuiltinRule:
-    def __init__(self, content: str, deps: list = None):
+    def __init__(self, content: str, deps: List[str]):
         self.content = content
-        self.deps = deps or []
+        self.deps = deps
 
 _up_to_15_digits = _build_repetition('[0-9]', 0, 15)
 
@@ -118,7 +118,7 @@ def __init__(self, *, prop_order, allow_fetch, dotall, raw_pattern):
 
     def _format_literal(self, literal):
         escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(
-            lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), literal
+            lambda m: GRAMMAR_LITERAL_ESCAPES[m.group(0)], literal
         )
         return f'"{escaped}"'
 
@@ -157,13 +157,13 @@ def _add_rule(self, name, rule):
         self._rules[key] = rule
         return key
 
-    def resolve_refs(self, schema: dict, url: str):
+    def resolve_refs(self, schema: Any, url: str):
         '''
             Resolves all $ref fields in the given schema, fetching any remote schemas,
             replacing $ref with absolute reference URL and populating self._refs with the
             respective referenced (sub)schema dictionaries.
         '''
-        def visit(n: dict):
+        def visit(n: Any):
             if isinstance(n, list):
                 return [visit(x) for x in n]
             elif isinstance(n, dict):
@@ -223,7 +223,7 @@ def _visit_pattern(self, pattern, name):
 
         assert pattern.startswith('^') and pattern.endswith('$'), 'Pattern must start with "^" and end with "$"'
         pattern = pattern[1:-1]
-        sub_rule_ids = {}
+        sub_rule_ids: Dict[str, str] = {}
 
         i = 0
         length = len(pattern)

examples/openai/api.py

@@ -1,5 +1,5 @@
 from abc import ABC
-from typing import Any, Dict, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 from pydantic import BaseModel, Json, TypeAdapter
 
 class FunctionCall(BaseModel):
@@ -16,7 +16,7 @@ class Message(BaseModel):
     name: Optional[str] = None
     tool_call_id: Optional[str] = None
     content: Optional[str]
-    tool_calls: Optional[list[ToolCall]] = None
+    tool_calls: Optional[List[ToolCall]] = None
 
 class ToolFunction(BaseModel):
     name: str
@@ -29,7 +29,7 @@ class Tool(BaseModel):
 
 class ResponseFormat(BaseModel):
     type: Literal["json_object"]
-    schema: Optional[Dict] = None
+    schema: Optional[Json[Any]] = None  # type: ignore
 
 class LlamaCppParams(BaseModel):
     n_predict: Optional[int] = None
@@ -56,8 +56,8 @@ class LlamaCppParams(BaseModel):
 
 class ChatCompletionRequest(LlamaCppParams):
     model: str
-    tools: Optional[list[Tool]] = None
-    messages: list[Message] = None
+    tools: Optional[List[Tool]] = None
+    messages: Optional[List[Message]] = None
     prompt: Optional[str] = None
     response_format: Optional[ResponseFormat] = None
 
@@ -67,7 +67,7 @@ class ChatCompletionRequest(LlamaCppParams):
 class Choice(BaseModel):
     index: int
     message: Message
-    logprobs: Optional[Json] = None
+    logprobs: Optional[Json[Any]] = None
     finish_reason: Union[Literal["stop"], Literal["tool_calls"]]
 
 class Usage(BaseModel):
@@ -84,7 +84,7 @@ class ChatCompletionResponse(BaseModel):
     object: Literal["chat.completion"]
     created: int
     model: str
-    choices: list[Choice]
+    choices: List[Choice]
     usage: Usage
     system_fingerprint: str
     error: Optional[CompletionError] = None

examples/openai/llama_cpp_server_api.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Optional
 from pydantic import Json
 
 from examples.openai.api import LlamaCppParams
@@ -9,4 +9,4 @@ class LlamaCppServerCompletionRequest(LlamaCppParams):
     cache_prompt: Optional[bool] = None
 
     grammar: Optional[str] = None
-    json_schema: Optional[Json] = None
+    json_schema: Optional[Json[Any]] = None