src/utils.py - master - fatboyslava/ffs_assistant

src/utils.py
132 строки4 KB
Viacheslav Kovalev
update easy
01 июл 2026, 23:37
01 июл 2026, 23:37feb52fd
"""Utility functions for the file-first knowledge-base assistant."""

import json
from typing import Any

from gigachat.models import Function, FunctionParameters
from pydantic import BaseModel

from src.config import DATA_DIR, DATA_ROOT

MAX_DATA_ELEMENTS = 50


def build_data_index() -> str:
    """Build a compact index of files and folders directly in data/ (first level only)."""

    elements: list[dict[str, str]] = []

    try:
        items = sorted(DATA_ROOT.iterdir(), key=lambda p: (not p.is_dir(), p.name))
    except PermissionError:
        return json.dumps([], ensure_ascii=False, indent=None)

    for item in items:
        if item.name.startswith("."):
            continue
        elements.append(
            {
                "path": item.name,
                "type": "directory" if item.is_dir() else "file",
            }
        )

    return json.dumps(elements, ensure_ascii=False, indent=None)


def validate_data_size() -> None:
    """Validate that data directory has at most MAX_DATA_ELEMENTS elements."""
    count = 0
    try:
        for item in DATA_ROOT.iterdir():
            if item.name.startswith("."):
                continue
            count += 1
            if count > MAX_DATA_ELEMENTS:
                break
    except PermissionError:
        pass

    if count > MAX_DATA_ELEMENTS:
        raise RuntimeError(
            f"В директории {DATA_DIR}/ слишком много элементов: {count}. "
            f"Максимально допустимо: {MAX_DATA_ELEMENTS}. "
            f"Сгруппируй файлы по разделам или уменьши количество документов."
        )


def _object_schema(model: type[BaseModel]) -> tuple[dict[str, Any], list[str]]:
    schema = model.model_json_schema()
    properties = {
        name: normalize_json_schema(property_schema)
        for name, property_schema in schema.get("properties", {}).items()
    }
    return properties, schema.get("required", [])


def normalize_json_schema(schema: dict[str, Any]) -> dict[str, Any]:
    """Convert Pydantic JSON Schema to the simpler shape accepted by GigaChat."""
    schema = dict(schema)

    if "anyOf" in schema:
        variants = [
            variant for variant in schema["anyOf"] if variant.get("type") != "null"
        ]
        if len(variants) == 1:
            variant = normalize_json_schema(variants[0])
            description = schema.get("description")
            schema = {**variant}
            if description:
                schema["description"] = description

    cleaned: dict[str, Any] = {}
    for key, value in schema.items():
        if key in {"anyOf", "default", "title"}:
            continue
        if key == "properties":
            cleaned[key] = {
                name: normalize_json_schema(property_schema)
                for name, property_schema in value.items()
            }
        elif key == "items" and isinstance(value, dict):
            cleaned[key] = normalize_json_schema(value)
        else:
            cleaned[key] = value

    return cleaned


def build_function(
    name: str,
    description: str,
    input_model: type[BaseModel],
    output_model: type[BaseModel],
    few_shot_examples: list[Any] | None = None,
) -> Function:
    """Build a GigaChat Function object from Pydantic input/output models."""
    input_properties, input_required = _object_schema(input_model)
    output_properties, output_required = _object_schema(output_model)

    return Function(
        name=name,
        description=description,
        parameters=FunctionParameters(
            type="object",
            properties=input_properties,
            required=input_required,
        ),
        return_parameters={
            "type": "object",
            "properties": output_properties,
            "required": output_required,
        },
        few_shot_examples=few_shot_examples,
    )


def build_functions() -> list[Function]:
    """Build the function descriptions passed to the model."""
    from src.tools.grep_tool.grep_tool import get_function as get_grep
    from src.tools.read_tool.read_tool import get_function as get_read

    return [get_grep(), get_read()]