API-key based backend using pydantic-ai + openai/litellm clients.
| 37 | logger = logging.getLogger(__name__) |
| 38 | |
| 39 | |
| 40 | def _run_usage(result: Any) -> dict[str, Any] | None: |
| 41 | """Token usage of a pydantic-ai run; ``usage`` is a property in pydantic-ai |
| 42 | 2.x and a method in earlier releases.""" |
| 43 | try: |
| 44 | usage = getattr(result, "usage", None) |
| 45 | if callable(usage): |
| 46 | usage = usage() |
| 47 | return usage_to_dict(usage) |
| 48 | except Exception: # noqa: BLE001 — usage is optional telemetry |
| 49 | return None |
| 50 | |
| 51 | |
| 52 | class PydanticAIBackend(LLMBackend): |
| 53 | """API-key based backend using pydantic-ai + openai/litellm clients.""" |
| 54 | |
| 55 | def __init__(self, config: Config) -> None: |
| 56 | self._config = config |
| 57 | self._fallback_models = create_fallback_models(config) |
| 58 | self._custom_instructions = config.get_prompt_addition() |
| 59 | self.last_usage: dict[str, Any] | None = None |
| 60 | |
| 61 | def complete( |
| 62 | self, |
| 63 | prompt: str, |
| 64 | *, |
| 65 | model: str | None = None, |
| 66 | ) -> str: |
| 67 | pop_last_usage() |
| 68 | result = call_llm(prompt, self._config, model=model) |
| 69 | self.last_usage = pop_last_usage() |
| 70 | return result |
| 71 | |
| 72 | async def run_update_agent( |
| 73 | self, |
| 74 | system_prompt: str, |
| 75 | user_prompt: str, |
| 76 | deps: CodeWikiDeps, |
| 77 | ) -> AgentReply: |
| 78 | agent = Agent( |
| 79 | self._fallback_models, |
| 80 | name=f"update:{deps.current_module_name}", |
| 81 | deps_type=CodeWikiDeps, |
| 82 | tools=[read_code_components_tool, str_replace_editor_tool], |
| 83 | system_prompt=system_prompt, |
| 84 | ) |
| 85 | started = time.time() |
| 86 | result = await agent.run(user_prompt, deps=deps) |
| 87 | seconds = time.time() - started |
| 88 | usage = _run_usage(result) |
| 89 | self.last_usage = usage |
| 90 | text = result.output if isinstance(result.output, str) else str(result.output) |
| 91 | return AgentReply(text=text, usage=usage, seconds=seconds) |
| 92 | |
| 93 | async def run_module_agent( |
| 94 | self, |
| 95 | module_name: str, |
| 96 | components: dict[str, Node], |