llm-codegen-research 2.10__tar.gz → 2.12__tar.gz

{llm_codegen_research-2.10 → llm_codegen_research-2.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llm-codegen-research
-Version: 2.10
+Version: 2.12
 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.10 → llm_codegen_research-2.12}/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.10 → llm_codegen_research-2.12}/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.12/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.10 → llm_codegen_research-2.12}/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.10 → llm_codegen_research-2.12}/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.10 → llm_codegen_research-2.12}/src/llm_cgr/llm/clients/mistral.py

@@ -3,7 +3,7 @@
 import os
 from typing import Any
 
-import mistralai
+from mistralai import client
 
 from llm_cgr.llm.clients.base import Base_LLM
 
@@ -31,7 +31,7 @@ class Mistral_LLM(Base_LLM):
             top_p=top_p,
             max_tokens=max_tokens,
         )
-        self._client = mistralai.Mistral(
+        self._client = client.Mistral(
             api_key=os.environ["MISTRAL_API_KEY"],
         )
 
@@ -71,8 +71,8 @@ class Mistral_LLM(Base_LLM):
         response = self._client.chat.complete(
             model=model,
             messages=input,
-            temperature=temperature if temperature is not None else mistralai.UNSET,
+            temperature=temperature if temperature is not None else client.UNSET,
             top_p=top_p,
-            max_tokens=max_tokens if max_tokens is not None else mistralai.UNSET,
+            max_tokens=max_tokens if max_tokens is not None else client.UNSET,
         )
         return response.choices[0].message.content

{llm_codegen_research-2.10 → llm_codegen_research-2.12}/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.10 → llm_codegen_research-2.12}/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.10 → llm_codegen_research-2.12}/src/llm_codegen_research.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llm-codegen-research
-Version: 2.10
+Version: 2.12
 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.10 → llm_codegen_research-2.12}/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

{llm_codegen_research-2.10 → llm_codegen_research-2.12}/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.10 → llm_codegen_research-2.12}/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.