llm-codegen-research 2.11__tar.gz → 2.13__tar.gz

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llm-codegen-research
-Version: 2.11
+Version: 2.13
 Summary: Useful classes and methods for researching code-generation by LLMs.
 Author-email: Lukas Twist <itsluketwist@gmail.com>
 Project-URL: Homepage, https://github.com/itsluketwist/llm-codegen-research

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/pyproject.toml

@@ -46,5 +46,11 @@ dev = [
     "uv",
 ]
 
+[tool.pytest.ini_options]
+markers = [
+    # tests that make real external api calls - excluded from ci runs
+    "api: marks tests as making external api calls",
+]
+
 [project.scripts]
 test_cuda = "llm_cgr.scripts.test_cuda:main"

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/__init__.py

@@ -45,7 +45,9 @@ try:
         GenerationProtocol,
         Mistral_LLM,
         OpenAI_LLM,
+        OpenAI_Tool_LLM,
         TogetherAI_LLM,
+        Tool,
         generate,
         generate_bool,
         generate_list,
@@ -64,6 +66,8 @@ try:
         "GenerationProtocol",
         "Mistral_LLM",
         "OpenAI_LLM",
+        "OpenAI_Tool_LLM",
+        "Tool",
         "TogetherAI_LLM",
         "generate",
         "generate_bool",

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/analyse/languages/__init__.py

@@ -1,6 +1,7 @@
 from llm_cgr.analyse.languages.code_data import CodeData
 from llm_cgr.analyse.languages.javascript import analyse_javascript_code
 from llm_cgr.analyse.languages.python import analyse_python_code
+from llm_cgr.analyse.languages.rust import analyse_rust_code
 
 
 def analyse_code(code: str, language: str | None) -> CodeData:
@@ -12,6 +13,8 @@ def analyse_code(code: str, language: str | None) -> CodeData:
             return analyse_python_code(code=code)
         elif language == "javascript":
             return analyse_javascript_code(code=code)
+        elif language == "rust":
+            return analyse_rust_code(code=code)
 
     except Exception as exc:
         return CodeData(

llm_codegen_research-2.13/src/llm_cgr/analyse/languages/rust.py

@@ -0,0 +1,193 @@
+"""Utility functions for Rust code analysis."""
+
+import re
+
+from llm_cgr.analyse.languages.code_data import CodeData
+
+
+# rust's three standard library crates - everything else comes from crates.io
+RUST_STDLIB = frozenset({"std", "core", "alloc"})
+
+# matches extern crate declarations: extern crate serde;
+_EXTERN_CRATE_RE = re.compile(r"^\s*extern\s+crate\s+(\w+)", re.MULTILINE)
+
+
+def _strip_comments(code: str) -> str:
+    """Remove // line comments and /* */ block comments from code."""
+    # remove block comments first (they can span multiple lines)
+    code = re.sub(r"/\*.*?\*/", "", code, flags=re.DOTALL)
+    # remove line comments
+    code = re.sub(r"//[^\n]*", "", code)
+    return code
+
+
+def _extract_use_statements(code: str) -> list[str]:
+    """
+    Extract raw use paths from all use declarations in the code.
+
+    Uses a bracket-aware scanner rather than plain regex, because braced imports
+    like `use std::{io, fmt};` span multiple tokens and can be multi-line.
+    Returns a list of raw path strings, e.g. ["std::collections::HashMap", "tokio"].
+    """
+    paths = []
+    i = 0
+    while i < len(code):
+        # find the next 'use' keyword (word boundary ensures we skip 'reuse', etc.)
+        match = re.search(r"\buse\s+", code[i:])
+        if match is None:
+            break
+
+        start = i + match.end()
+        depth = 0
+        j = start
+
+        # scan forward until we hit a ';' at brace-depth 0
+        while j < len(code):
+            c = code[j]
+            if c == "{":
+                depth += 1
+            elif c == "}":
+                depth -= 1
+            elif c == ";" and depth == 0:
+                paths.append(code[start:j].strip())
+                break
+            j += 1
+
+        i = start  # advance past the 'use' keyword we just processed
+
+    return paths
+
+
+def _split_top_level(s: str) -> list[str]:
+    """
+    Split a comma-separated string into items, respecting brace nesting.
+
+    Used to split the contents of {...} groups in use paths.
+    """
+    items = []
+    depth = 0
+    current: list[str] = []
+
+    for c in s:
+        if c == "{":
+            depth += 1
+            current.append(c)
+        elif c == "}":
+            depth -= 1
+            current.append(c)
+        elif c == "," and depth == 0:
+            items.append("".join(current).strip())
+            current = []
+        else:
+            current.append(c)
+
+    if current:
+        items.append("".join(current).strip())
+
+    return items
+
+
+def _expand_use_path(path: str, prefix: str = "") -> list[str]:
+    """
+    Recursively expand a Rust use path into fully-qualified dotted import strings.
+
+    Handles simple paths, braced groups, wildcards, and 'as' aliases.
+    For example:
+      "std::collections::HashMap" -> ["std.collections.HashMap"]
+      "std::{io::Read, fmt}"      -> ["std.io.Read", "std.fmt"]
+      "std::io::{self, Write}"    -> ["std.io", "std.io.Write"]
+    Returns a list of dotted import paths.
+    """
+    # strip any trailing 'as alias' at this level of the path
+    path = re.sub(r"\s+as\s+\w+\s*$", "", path).strip()
+
+    if not path:
+        return []
+
+    # handle wildcard: use std::collections::*;
+    if path.endswith("::*"):
+        prefix_part = path[:-3].replace("::", ".")
+        full = f"{prefix}.{prefix_part}" if prefix else prefix_part
+        return [f"{full}.*"]
+
+    brace_idx = path.find("{")
+
+    if brace_idx == -1:
+        # simple path, no braces: std::collections::HashMap
+        converted = path.replace("::", ".")
+        full = f"{prefix}.{converted}" if prefix else converted
+        return [full]
+
+    # braced path: std::{io::Read, fmt} or std::io::{self, Write}
+    # everything before the brace is the common prefix
+    before_brace = path[:brace_idx].rstrip(":").replace("::", ".")
+    new_prefix = (
+        f"{prefix}.{before_brace}"
+        if (prefix and before_brace)
+        else (prefix or before_brace)
+    )
+
+    # extract the content inside the outermost braces
+    close_idx = path.rfind("}")
+    inner = path[brace_idx + 1 : close_idx]
+
+    results = []
+    for item in _split_top_level(inner):
+        item = item.strip()
+        if not item:
+            continue
+
+        if item == "self":
+            # 'self' refers to the module itself (the prefix path)
+            results.append(new_prefix)
+        else:
+            results.extend(_expand_use_path(item, prefix=new_prefix))
+
+    return results
+
+
+def analyse_rust_code(code: str) -> CodeData:
+    """
+    Analyse Rust code to extract imported crates and their paths.
+
+    Only extracts use and extern crate declarations; no usage tracking or syntax
+    validation is performed (valid is always True).
+    Returns a CodeData object with import information.
+    """
+    std_libs: set[str] = set()
+    ext_libs: set[str] = set()
+    imports: set[str] = set()
+
+    # strip comments so 'use' inside comments isn't matched
+    clean_code = _strip_comments(code)
+
+    # process each use declaration
+    for raw_path in _extract_use_statements(clean_code):
+        for dotted_path in _expand_use_path(raw_path):
+            imports.add(dotted_path)
+            # the top-level segment is the crate name
+            top_level = dotted_path.split(".")[0]
+            if top_level in RUST_STDLIB:
+                std_libs.add(top_level)
+            else:
+                ext_libs.add(top_level)
+
+    # handle extern crate declarations (older rust style, still used occasionally)
+    for match in _EXTERN_CRATE_RE.finditer(clean_code):
+        crate_name = match.group(1)
+        # 'self' and 'std' as extern crate are special cases, skip them
+        if crate_name in ("self", "super"):
+            continue
+        imports.add(crate_name)
+        if crate_name in RUST_STDLIB:
+            std_libs.add(crate_name)
+        else:
+            ext_libs.add(crate_name)
+
+    return CodeData(
+        valid=True,
+        std_libs=std_libs,
+        ext_libs=ext_libs,
+        imports=imports,
+        lib_usage={},
+    )

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/llm/__init__.py

@@ -5,7 +5,9 @@ from llm_cgr.llm.clients import (
     GenerationProtocol,
     Mistral_LLM,
     OpenAI_LLM,
+    OpenAI_Tool_LLM,
     TogetherAI_LLM,
+    Tool,
     get_llm,
 )
 from llm_cgr.llm.generate import generate, generate_bool, generate_list
@@ -24,6 +26,8 @@ __all__ = [
     "GenerationProtocol",
     "Mistral_LLM",
     "OpenAI_LLM",
+    "OpenAI_Tool_LLM",
+    "Tool",
     "TogetherAI_LLM",
     "get_llm",
     "generate",

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/llm/clients/__init__.py

@@ -6,6 +6,7 @@ from llm_cgr.llm.clients.deepseek import DeepSeek_LLM
 from llm_cgr.llm.clients.mistral import Mistral_LLM
 from llm_cgr.llm.clients.nscale import Nscale_LLM
 from llm_cgr.llm.clients.openai import OpenAI_LLM
+from llm_cgr.llm.clients.openai_tool import OpenAI_Tool_LLM, Tool
 from llm_cgr.llm.clients.protocol import GenerationProtocol
 from llm_cgr.llm.clients.together import TogetherAI_LLM
 
@@ -27,9 +28,13 @@ def get_llm(
     top_p: float | None = None,
     max_tokens: int | None = None,
     provider: str | None = None,
+    tools: list[Tool] | None = None,
 ) -> GenerationProtocol:
     """
     Initialise the correct LLM client for the given model.
+
+    If tools are provided, returns an OpenAI_Tool_LLM instance. Tool calls
+    are currently only supported for OpenAI models.
     """
     llm_class: type[Base_LLM]
     if provider is not None:
@@ -45,6 +50,21 @@ def get_llm(
     else:
         llm_class = TogetherAI_LLM
 
+    # if tools are requested, use the tool-enabled subclass (openai only for now)
+    if tools is not None:
+        if llm_class is not OpenAI_LLM:
+            raise NotImplementedError(
+                "Tool calls are only supported for OpenAI models."
+            )
+        return OpenAI_Tool_LLM(
+            tools=tools,
+            model=model,
+            system=system,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+        )
+
     return llm_class(
         model=model,
         system=system,
@@ -60,7 +80,9 @@ __all__ = [
     "DeepSeek_LLM",
     "GenerationProtocol",
     "OpenAI_LLM",
+    "OpenAI_Tool_LLM",
     "TogetherAI_LLM",
     "Mistral_LLM",
+    "Tool",
     "get_llm",
 ]

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/llm/clients/anthropic.py

@@ -1,8 +1,9 @@
 """Classes to access LLMs via the Anthropic Claude API."""
 
-from typing import Any
+from typing import Any, cast
 
 import anthropic
+from anthropic.types import MessageParam, TextBlock
 
 from llm_cgr.defaults import DEFAULT_MAX_TOKENS
 from llm_cgr.llm.clients.base import Base_LLM
@@ -69,10 +70,11 @@ class Anthropic_LLM(Base_LLM):
         """Generate a model response from the Anthropic API."""
         response = self._client.messages.create(
             model=model,
-            system=system or self._system or anthropic.NOT_GIVEN,
-            messages=input,
-            temperature=temperature if temperature is not None else anthropic.NOT_GIVEN,
-            top_p=top_p if top_p is not None else anthropic.NOT_GIVEN,
+            system=system or self._system or anthropic.omit,
+            messages=cast(list[MessageParam], input),
+            temperature=temperature if temperature is not None else anthropic.omit,
+            top_p=top_p if top_p is not None else anthropic.omit,
             max_tokens=max_tokens if max_tokens is not None else DEFAULT_MAX_TOKENS,
         )
-        return response.content[0].text
+        # cast to TextBlock as non-tool, non-thinking requests always return text
+        return cast(TextBlock, response.content[0]).text

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/llm/clients/deepseek.py

@@ -1,9 +1,10 @@
 """Class to access LLMs via the OpenAI API."""
 
 import os
-from typing import Any
+from typing import Any, cast
 
 import openai
+from openai.types.chat import ChatCompletionMessageParam
 
 from llm_cgr.llm.clients.base import Base_LLM
 
@@ -67,10 +68,11 @@ class DeepSeek_LLM(Base_LLM):
     ) -> str:
         """Generate a model response from the OpenAI API."""
         response = self._client.chat.completions.create(
-            messages=input,
+            messages=cast(list[ChatCompletionMessageParam], input),
             model=model,
             temperature=temperature if temperature is not None else openai.omit,
             top_p=top_p if top_p is not None else openai.omit,
             max_completion_tokens=max_tokens if max_tokens is not None else openai.omit,
         )
-        return response.choices[0].message.content
+        # cast to str as text completions always return string content
+        return cast(str, response.choices[0].message.content)

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/llm/clients/nscale.py

@@ -1,9 +1,10 @@
 """Class to access LLMs via the OpenAI API."""
 
 import os
-from typing import Any
+from typing import Any, cast
 
 import openai
+from openai.types.chat import ChatCompletionMessageParam
 
 from llm_cgr.llm.clients.base import Base_LLM
 
@@ -67,10 +68,11 @@ class Nscale_LLM(Base_LLM):
     ) -> str:
         """Generate a model response from the OpenAI API."""
         response = self._client.chat.completions.create(
-            messages=input,
+            messages=cast(list[ChatCompletionMessageParam], input),
             model=model,
             temperature=temperature if temperature is not None else openai.omit,
             top_p=top_p if top_p is not None else openai.omit,
             max_completion_tokens=max_tokens if max_tokens is not None else openai.omit,
         )
-        return response.choices[0].message.content
+        # cast to str as text completions always return string content
+        return cast(str, response.choices[0].message.content)

llm_codegen_research-2.13/src/llm_cgr/llm/clients/openai_tool.py

@@ -0,0 +1,258 @@
+"""OpenAI client subclass with an agentic tool-call loop."""
+
+import json
+from dataclasses import dataclass
+from typing import Any, Callable, cast
+
+import openai
+from openai.types.responses import ResponseFunctionToolCall, ResponseInputItemParam
+
+from llm_cgr.llm.clients.openai import OpenAI_LLM
+
+
+# maximum number of tool-call iterations per request, to prevent runaway loops
+MAX_TOOL_ITERATIONS: int = 10
+
+
+@dataclass
+class Tool:
+    """
+    A tool (function) that the model can call during generation.
+
+    Attributes:
+        name: The function name the model uses to call this tool.
+        description: Describes what the tool does; the model uses this
+            to decide when to call it.
+        parameters: A JSON schema dict describing the function's parameters.
+        fn: The Python callable to invoke; must accept kwargs matching the
+            schema and return a str result.
+    """
+
+    name: str
+    description: str
+    parameters: dict[str, Any]
+    fn: Callable[..., str]
+
+
+class OpenAI_Tool_LLM(OpenAI_LLM):
+    """OpenAI client with an agentic tool-call loop.
+
+    Tools are supplied at construction time and used for all subsequent
+    generate() and chat() calls. The client handles the full loop internally:
+    call the API, execute any tool calls, feed results back, repeat until the
+    model produces a final text response.
+    """
+
+    def __init__(
+        self,
+        tools: list[Tool],
+        model: str | None = None,
+        system: str | None = None,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+    ) -> None:
+        """
+        Initialise the OpenAI tool client.
+
+        Requires the OPENAI_API_KEY environment variable to be set.
+        """
+        super().__init__(
+            model=model,
+            system=system,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+        )
+        self._tools = tools
+        # cumulative count of individual tool calls made by this instance
+        self._tool_calls: int = 0
+
+    @property
+    def tool_calls(self) -> int:
+        """Total number of tool calls made by this client since instantiation.
+
+        Returns the cumulative count across all generate() and chat() calls.
+        Tip: record the value before a call and subtract to get the count for
+        that specific call.
+        """
+        return self._tool_calls
+
+    def _build_tool_param(
+        self,
+        tool: Tool,
+    ) -> dict[str, Any]:
+        """Convert a Tool dataclass to the dict format the OpenAI Responses API expects."""
+        return {
+            "type": "function",
+            "name": tool.name,
+            "description": tool.description,
+            "parameters": tool.parameters,
+        }
+
+    def _run_tool_loop(
+        self,
+        messages: list[dict[str, Any]],
+        model: str,
+        temperature: float | None,
+        top_p: float | None,
+        max_tokens: int | None,
+    ) -> str:
+        """Run the agentic tool-call loop for a single turn.
+
+        Calls the OpenAI API in a loop, executing any tool calls the model
+        requests, until the model produces a final text response or the
+        MAX_TOOL_ITERATIONS safety limit is reached.
+
+        Returns the final text response.
+        """
+        # convert Tool dataclasses to the API's function-tool format
+        api_tools = [self._build_tool_param(t) for t in self._tools]
+
+        # build a name -> Tool lookup map for fast dispatch during the loop
+        tool_map = {t.name: t for t in self._tools}
+
+        # shallow copy so intermediate tool-call scaffolding never mutates the
+        # caller's message list (prevents corruption of the chat history).
+        # typed as list[Any] so we can freely append both plain message dicts
+        # and the richer tool-call dicts without fighting the type checker.
+        current_input: list[Any] = list(messages)
+
+        for _ in range(MAX_TOOL_ITERATIONS):
+            response = self._client.responses.create(
+                input=cast(list[ResponseInputItemParam], current_input),
+                model=model,
+                temperature=temperature if temperature is not None else openai.omit,
+                top_p=top_p if top_p is not None else openai.omit,
+                max_output_tokens=max_tokens if max_tokens is not None else openai.omit,
+                tools=cast(Any, api_tools),
+            )
+
+            # collect any function calls the model requested in this response
+            function_calls = [
+                item for item in response.output if item.type == "function_call"
+            ]
+
+            # no tool calls means the model has produced its final text answer
+            if not function_calls:
+                return response.output_text
+
+            # increment the cumulative counter; parallel calls count individually
+            self._tool_calls += len(function_calls)
+
+            # process each tool call: the OpenAI Responses API requires that the
+            # function_call item appears in the next input before its matching
+            # function_call_output item
+            for _call in function_calls:
+                # cast to the concrete type so we can access .call_id/.name/.arguments
+                call = cast(ResponseFunctionToolCall, _call)
+
+                # append the function_call itself so the model sees what it called
+                current_input.append(
+                    {
+                        "type": "function_call",
+                        "call_id": call.call_id,
+                        "name": call.name,
+                        "arguments": call.arguments,
+                    }
+                )
+
+                # deserialise the model's json argument string and call the local fn
+                kwargs = json.loads(call.arguments)
+                result = tool_map[call.name].fn(**kwargs)
+
+                # append the result so the model can read it on the next turn
+                current_input.append(
+                    {
+                        "type": "function_call_output",
+                        "call_id": call.call_id,
+                        "output": result,
+                    }
+                )
+
+            # loop continues: enriched input is sent back to the model
+
+        # safety fallback: return whatever text the model produced on the last turn
+        return response.output_text
+
+    def generate(
+        self,
+        user: str,
+        system: str | None = None,
+        model: str | None = None,
+        samples: int = 1,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+    ) -> list[str]:
+        """Generate model responses via the agentic tool-call loop."""
+        _model = model or self._model
+        if _model is None:
+            raise ValueError("Model must be specified for LLM APIs.")
+
+        messages = self._build_input(
+            user=user,
+            system=system or self._system,
+        )
+
+        _generations = []
+        for _ in range(samples):
+            result = self._run_tool_loop(
+                messages=messages,
+                model=_model,
+                temperature=temperature or self._temperature,
+                top_p=top_p or self._top_p,
+                max_tokens=max_tokens or self._max_tokens,
+            )
+            _generations.append(result)
+
+        return _generations
+
+    def chat(
+        self,
+        user: str,
+        system: str | None = None,
+        model: str | None = None,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+    ) -> str:
+        """Run a chat turn via the agentic tool-call loop.
+
+        Manages self._history identically to the base class — only the final
+        text response is appended, not intermediate tool-call scaffolding.
+        """
+        _model = model or self._model
+        if _model is None:
+            raise ValueError("Model must be specified for LLM APIs.")
+
+        if self._history is None:
+            self._history = self._build_input(
+                user=user,
+                system=system or self._system,
+            )
+        else:
+            self._history.append(
+                self._build_message(
+                    role="user",
+                    content=user,
+                )
+            )
+
+        # _run_tool_loop operates on a shallow copy of self._history, so
+        # intermediate tool-call items never appear in the chat history
+        response = self._run_tool_loop(
+            messages=self._history,
+            model=_model,
+            temperature=temperature or self._temperature,
+            top_p=top_p or self._top_p,
+            max_tokens=max_tokens or self._max_tokens,
+        )
+
+        self._history.append(
+            self._build_message(
+                role="assistant",
+                content=response,
+            )
+        )
+        return response

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_cgr/llm/clients/together.py

@@ -1,6 +1,6 @@
 """Class to access LLMs via the TogetherAI API."""
 
-from typing import Any
+from typing import Any, cast
 
 import together
 
@@ -64,9 +64,12 @@ class TogetherAI_LLM(Base_LLM):
         """Generate a model response from the TogetherAI API."""
         response = self._client.chat.completions.create(
             model=model,
-            messages=input,
-            temperature=temperature,
-            top_p=top_p,
-            max_tokens=max_tokens,
+            messages=cast(Any, input),
+            temperature=temperature if temperature is not None else together.omit,
+            top_p=top_p if top_p is not None else together.omit,
+            max_tokens=max_tokens if max_tokens is not None else together.omit,
         )
-        return response.choices[0].message.content
+        # cast to Any first as together doesn't publicly export the message type,
+        # then cast content to str as text completions always have it set
+        message = cast(Any, response.choices[0].message)
+        return cast(str, message.content)

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_codegen_research.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llm-codegen-research
-Version: 2.11
+Version: 2.13
 Summary: Useful classes and methods for researching code-generation by LLMs.
 Author-email: Lukas Twist <itsluketwist@gmail.com>
 Project-URL: Homepage, https://github.com/itsluketwist/llm-codegen-research

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/src/llm_codegen_research.egg-info/SOURCES.txt

@@ -15,6 +15,7 @@ src/llm_cgr/analyse/languages/__init__.py
 src/llm_cgr/analyse/languages/code_data.py
 src/llm_cgr/analyse/languages/javascript.py
 src/llm_cgr/analyse/languages/python.py
+src/llm_cgr/analyse/languages/rust.py
 src/llm_cgr/llm/__init__.py
 src/llm_cgr/llm/generate.py
 src/llm_cgr/llm/prompts.py
@@ -25,6 +26,7 @@ src/llm_cgr/llm/clients/deepseek.py
 src/llm_cgr/llm/clients/mistral.py
 src/llm_cgr/llm/clients/nscale.py
 src/llm_cgr/llm/clients/openai.py
+src/llm_cgr/llm/clients/openai_tool.py
 src/llm_cgr/llm/clients/protocol.py
 src/llm_cgr/llm/clients/together.py
 src/llm_cgr/scripts/test_cuda.py
@@ -38,4 +40,5 @@ tests/test_enums.py
 tests/test_json_utils.py
 tests/test_llm_api.py
 tests/test_llm_local.py
+tests/test_llm_tool.py
 tests/test_utils.py

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/tests/test_enums.py

@@ -1,4 +1,5 @@
 from enum import auto
+from typing import Any
 
 from llm_cgr import OptionsEnum
 
@@ -26,8 +27,8 @@ def test_options_enum():
     assert (TestEnum.ONE != "One") is False
     assert (TestEnum.ONE != "ONE") is False
 
-    # check that the enum can be hashed
-    test_dict = {TestEnum.TWO: "value"}
+    # check that the enum can be hashed and string keys match enum keys at runtime
+    test_dict: dict[Any, str] = {TestEnum.TWO: "value"}
     assert test_dict[TestEnum.TWO] == "value"
     assert test_dict["two"] == "value"
     assert {TestEnum.THREE} == {"three"}

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/tests/test_llm_api.py

@@ -1,8 +1,14 @@
 """Test our connection and usage of the LLM APIs."""
 
+import pytest
+
 from llm_cgr import BASE_SYSTEM_PROMPT, generate, generate_bool, generate_list, get_llm
 
 
+# mark all tests in this file as api tests, so they can be excluded in ci
+pytestmark = pytest.mark.api
+
+
 def test_generate(model):
     """
     Test the generate method.

{llm_codegen_research-2.11 → llm_codegen_research-2.13}/tests/test_llm_local.py

@@ -10,8 +10,10 @@ from llm_cgr import (
     Mistral_LLM,
     OpenAI_LLM,
     TogetherAI_LLM,
+    Tool,
     generate_bool,
     generate_list,
+    get_llm,
 )
 
 
@@ -129,6 +131,30 @@ def test_build_input():
     ]
 
 
+def test_tools_unsupported_provider():
+    """
+    Test that passing tools to a non-OpenAI provider raises NotImplementedError.
+
+    No API call is made because the error fires inside get_llm() before any
+    network request.
+    """
+    dummy_tool = Tool(
+        name="dummy",
+        description="A dummy tool.",
+        parameters={"type": "object", "properties": {}},
+        fn=lambda: "result",
+    )
+
+    # anthropic is not yet supported for tool calls
+    with pytest.raises(
+        NotImplementedError, match="Tool calls are only supported for OpenAI models."
+    ):
+        get_llm(
+            model="claude-3-5-haiku-20241022",
+            tools=[dummy_tool],
+        )
+
+
 @pytest.mark.parametrize(
     "response,error",
     [

llm_codegen_research-2.13/tests/test_llm_tool.py

@@ -0,0 +1,96 @@
+"""Test the OpenAI_Tool_LLM agentic tool-call loop."""
+
+import pytest
+
+from llm_cgr import OpenAI_Tool_LLM, Tool
+
+
+# mark all tests in this file as api tests, so they can be excluded in ci
+pytestmark = pytest.mark.api
+
+
+def test_tool_call_generate(openai_model):
+    """
+    Test that the OpenAI tool client runs the agentic loop and returns the
+    correct answer via the tool.
+
+    Uses an addition tool: the model must call it to get the answer, so we can
+    verify a real tool call happened (the model cannot guess what our local
+    function returns without calling it).
+    """
+
+    def add(a: int, b: int) -> str:
+        """Add two integers and return the result as a string."""
+        return str(a + b)
+
+    add_tool = Tool(
+        name="add",
+        description="Add two integers together and return the result.",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "integer", "description": "The first integer."},
+                "b": {"type": "integer", "description": "The second integer."},
+            },
+            "required": ["a", "b"],
+            "additionalProperties": False,
+        },
+        fn=add,
+    )
+
+    llm = OpenAI_Tool_LLM(tools=[add_tool], model=openai_model)
+    responses = llm.generate(
+        user="Use the add tool to compute 3 + 4. What is the result?"
+    )
+
+    assert isinstance(responses, list)
+    assert len(responses) == 1
+
+    result = responses[0]
+    assert isinstance(result, str)
+    assert len(result) > 0
+    # the correct sum proves the tool was actually called
+    assert "7" in result
+    assert llm.tool_calls >= 1
+
+
+def test_tool_call_chat(openai_model):
+    """
+    Test that tool calls work correctly in a chat session, and that
+    intermediate tool-call scaffolding does not corrupt the chat history.
+    """
+
+    def multiply(a: int, b: int) -> str:
+        """Multiply two integers and return the result as a string."""
+        return str(a * b)
+
+    multiply_tool = Tool(
+        name="multiply",
+        description="Multiply two integers together and return the result.",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "integer", "description": "The first integer."},
+                "b": {"type": "integer", "description": "The second integer."},
+            },
+            "required": ["a", "b"],
+            "additionalProperties": False,
+        },
+        fn=multiply,
+    )
+
+    llm = OpenAI_Tool_LLM(tools=[multiply_tool], model=openai_model)
+    response = llm.chat(
+        user="Use the multiply tool to compute 6 * 7. What is the result?"
+    )
+
+    assert isinstance(response, str)
+    assert "42" in response
+
+    # history should only contain the user turn and the final assistant response;
+    # no intermediate function_call or function_call_output items should have leaked in
+    history = llm.history
+    assert len(history) == 2
+    assert history[0]["role"] == "user"
+    assert history[1]["role"] == "assistant"
+    assert llm.tool_calls >= 1