优化目录结构

This commit is contained in:
2026-07-03 20:49:36 +08:00
parent 04fefcf488
commit 8073dd63a7
10 changed files with 220 additions and 238 deletions
+9 -1
View File
@@ -2,6 +2,15 @@
一个使用 FastAPI 和 DeepSeek API 的多轮 Agent 服务。每个会话的系统提示词、消息历史和工具调用过程保存在本地 JSON 文件中。 一个使用 FastAPI 和 DeepSeek API 的多轮 Agent 服务。每个会话的系统提示词、消息历史和工具调用过程保存在本地 JSON 文件中。
## 代码结构
- `routes.py`:集中定义所有 HTTP 路由。
- `schemas.py`:集中定义 API 请求和响应结构。
- `domain.py`:定义会话、消息和工具调用的持久化模型。
- `tools.py`:集中定义工具注册表、参数结构和工具函数。
- `service.py`:处理对话、Agent 循环和业务规则。
- `app.py`:创建 FastAPI 应用并管理生命周期。
## 启动 ## 启动
```bash ```bash
@@ -38,7 +47,6 @@ API 返回及存储的每条消息都包含 UTC `created_at`。只有模型完
模型可以按需调用服务端白名单工具: 模型可以按需调用服务端白名单工具:
- `get_current_time`:查询指定 IANA 时区的当前时间。 - `get_current_time`:查询指定 IANA 时区的当前时间。
- `calculate`:计算受限的基础算术表达式,不执行任意 Python 代码。
工具调用无需增加请求参数。服务端会执行工具并把结果返回模型,直到模型生成最终回答。发送消息接口通过 `tools_use` 返回本轮使用的工具名称,完整调用过程可通过历史接口查询;未调用工具时 `tools_use` 为空数组。 工具调用无需增加请求参数。服务端会执行工具并把结果返回模型,直到模型生成最终回答。发送消息接口通过 `tools_use` 返回本轮使用的工具名称,完整调用过程可通过历史接口查询;未调用工具时 `tools_use` 为空数组。
+10 -69
View File
@@ -1,19 +1,13 @@
from contextlib import asynccontextmanager
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Request, status from fastapi import FastAPI
from openai import AsyncOpenAI from openai import AsyncOpenAI
from config import Settings from config import Settings
from models import ( from routes import router
CreateSessionRequest, from service import ChatService
CreateSessionResponse, from storage import JsonSessionStorage
SendMessageRequest,
SendMessageResponse,
SessionHistoryResponse,
)
from service import ChatProviderError, ChatService
from storage import JsonSessionStorage, SessionNotFoundError, SessionStorageError
from tools import ToolRegistry from tools import ToolRegistry
@@ -22,7 +16,6 @@ def create_app(
client: AsyncOpenAI | None = None, client: AsyncOpenAI | None = None,
) -> FastAPI: ) -> FastAPI:
resolved_settings = settings or Settings.from_env() resolved_settings = settings or Settings.from_env()
storage = JsonSessionStorage(resolved_settings.data_dir)
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
@@ -34,75 +27,23 @@ def create_app(
api_key=resolved_settings.api_key, api_key=resolved_settings.api_key,
base_url=resolved_settings.base_url, base_url=resolved_settings.base_url,
) )
app.state.chat_service = ChatService( app.state.chat_service = ChatService(
storage=storage, storage=JsonSessionStorage(resolved_settings.data_dir),
client=resolved_client, client=resolved_client,
model=resolved_settings.model, model=resolved_settings.model,
default_system_prompt=resolved_settings.default_system_prompt,
tool_registry=ToolRegistry(resolved_settings.tool_timeout_seconds), tool_registry=ToolRegistry(resolved_settings.tool_timeout_seconds),
max_tool_rounds=resolved_settings.max_tool_rounds, max_tool_rounds=resolved_settings.max_tool_rounds,
max_tool_calls_per_turn=resolved_settings.max_tool_calls_per_turn, max_tool_calls_per_turn=resolved_settings.max_tool_calls_per_turn,
) )
yield yield
if client is None: if client is None:
await resolved_client.close() await resolved_client.close()
app = FastAPI(title="Simple Chat API", version="0.1.0", lifespan=lifespan) app = FastAPI(title="Simple Chat API", version="0.1.0", lifespan=lifespan)
app.include_router(router)
@app.post(
"/sessions",
response_model=CreateSessionResponse,
status_code=status.HTTP_201_CREATED,
)
async def create_session(
body: CreateSessionRequest | None = None,
) -> CreateSessionResponse:
system_prompt = (
body.system_prompt
if body is not None and body.system_prompt is not None
else resolved_settings.default_system_prompt
)
try:
session = await storage.create(system_prompt)
except SessionStorageError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
return CreateSessionResponse(
session_id=session.session_id,
system_prompt=session.system_prompt,
created_at=session.created_at,
)
@app.post(
"/sessions/{session_id}/messages",
response_model=SendMessageResponse,
)
async def send_message(
session_id: str, body: SendMessageRequest, request: Request
) -> SendMessageResponse:
service: ChatService = request.app.state.chat_service
try:
return await service.generate_response(session_id, body.content)
except SessionNotFoundError as exc:
raise HTTPException(status_code=404, detail="session not found") from exc
except ChatProviderError as exc:
raise HTTPException(status_code=502, detail=str(exc)) from exc
except SessionStorageError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.get(
"/sessions/{session_id}/messages",
response_model=SessionHistoryResponse,
)
async def get_session_history(
session_id: str, request: Request
) -> SessionHistoryResponse:
service: ChatService = request.app.state.chat_service
try:
return await service.get_session_history(session_id)
except SessionNotFoundError as exc:
raise HTTPException(status_code=404, detail="session not found") from exc
except SessionStorageError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
return app return app
+3 -60
View File
@@ -1,7 +1,7 @@
from datetime import datetime from datetime import datetime
from typing import Annotated, Literal from typing import Annotated, Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from pydantic import BaseModel, ConfigDict, Field, model_validator
class ToolFunctionCall(BaseModel): class ToolFunctionCall(BaseModel):
@@ -41,6 +41,8 @@ Message = Annotated[
UserMessage | AssistantMessage | ToolMessage, UserMessage | AssistantMessage | ToolMessage,
Field(discriminator="role"), Field(discriminator="role"),
] ]
class Session(BaseModel): class Session(BaseModel):
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
@@ -63,62 +65,3 @@ class Session(BaseModel):
if isinstance(message, dict): if isinstance(message, dict):
message.setdefault("created_at", fallback) message.setdefault("created_at", fallback)
return data return data
class CreateSessionRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
system_prompt: str | None = None
@field_validator("system_prompt")
@classmethod
def validate_system_prompt(cls, value: str | None) -> str | None:
if value is not None and not value.strip():
raise ValueError("system_prompt must not be blank")
return value
class CreateSessionResponse(BaseModel):
session_id: str
system_prompt: str
created_at: datetime
class SendMessageRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
content: str
@field_validator("content")
@classmethod
def validate_content(cls, value: str) -> str:
if not value.strip():
raise ValueError("content must not be blank")
return value
class TokenUsage(BaseModel):
prompt_tokens: int
completion_tokens: int
total_tokens: int
class FinalAssistantMessage(BaseModel):
role: Literal["assistant"] = "assistant"
content: str
created_at: datetime
class SendMessageResponse(BaseModel):
session_id: str
message: FinalAssistantMessage
tools_use: list[str] = Field(default_factory=list)
usage: TokenUsage | None
class SessionHistoryResponse(BaseModel):
session_id: str
system_prompt: str
created_at: datetime
updated_at: datetime
messages: list[Message]
+66
View File
@@ -0,0 +1,66 @@
from fastapi import APIRouter, HTTPException, Request, status
from schemas import (
CreateSessionRequest,
CreateSessionResponse,
SendMessageRequest,
SendMessageResponse,
SessionHistoryResponse,
)
from service import ChatProviderError, ChatService
from storage import SessionNotFoundError, SessionStorageError
router = APIRouter(prefix="/sessions", tags=["sessions"])
def get_chat_service(request: Request) -> ChatService:
return request.app.state.chat_service
@router.post(
"",
response_model=CreateSessionResponse,
status_code=status.HTTP_201_CREATED,
)
async def create_session(
request: Request,
body: CreateSessionRequest | None = None,
) -> CreateSessionResponse:
try:
return await get_chat_service(request).create_session(
body.system_prompt if body is not None else None
)
except SessionStorageError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@router.post("/{session_id}/messages", response_model=SendMessageResponse)
async def send_message(
session_id: str,
body: SendMessageRequest,
request: Request,
) -> SendMessageResponse:
try:
return await get_chat_service(request).generate_response(
session_id, body.content
)
except SessionNotFoundError as exc:
raise HTTPException(status_code=404, detail="session not found") from exc
except ChatProviderError as exc:
raise HTTPException(status_code=502, detail=str(exc)) from exc
except SessionStorageError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@router.get("/{session_id}/messages", response_model=SessionHistoryResponse)
async def get_session_history(
session_id: str,
request: Request,
) -> SessionHistoryResponse:
try:
return await get_chat_service(request).get_session_history(session_id)
except SessionNotFoundError as exc:
raise HTTPException(status_code=404, detail="session not found") from exc
except SessionStorageError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
+65
View File
@@ -0,0 +1,65 @@
from datetime import datetime
from typing import Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator
from domain import Message
class CreateSessionRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
system_prompt: str | None = None
@field_validator("system_prompt")
@classmethod
def validate_system_prompt(cls, value: str | None) -> str | None:
if value is not None and not value.strip():
raise ValueError("system_prompt must not be blank")
return value
class CreateSessionResponse(BaseModel):
session_id: str
system_prompt: str
created_at: datetime
class SendMessageRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
content: str
@field_validator("content")
@classmethod
def validate_content(cls, value: str) -> str:
if not value.strip():
raise ValueError("content must not be blank")
return value
class TokenUsage(BaseModel):
prompt_tokens: int
completion_tokens: int
total_tokens: int
class FinalAssistantMessage(BaseModel):
role: Literal["assistant"] = "assistant"
content: str
created_at: datetime
class SendMessageResponse(BaseModel):
session_id: str
message: FinalAssistantMessage
tools_use: list[str] = Field(default_factory=list)
usage: TokenUsage | None
class SessionHistoryResponse(BaseModel):
session_id: str
system_prompt: str
created_at: datetime
updated_at: datetime
messages: list[Message]
+22 -5
View File
@@ -4,19 +4,22 @@ from typing import Any
from openai import APIError, AsyncOpenAI from openai import APIError, AsyncOpenAI
from models import ( from domain import (
AssistantMessage, AssistantMessage,
FinalAssistantMessage,
Message, Message,
SendMessageResponse,
Session, Session,
SessionHistoryResponse,
TokenUsage,
ToolCall, ToolCall,
ToolFunctionCall, ToolFunctionCall,
ToolMessage, ToolMessage,
UserMessage, UserMessage,
) )
from schemas import (
CreateSessionResponse,
FinalAssistantMessage,
SendMessageResponse,
SessionHistoryResponse,
TokenUsage,
)
from storage import JsonSessionStorage from storage import JsonSessionStorage
from tools import ToolRegistry from tools import ToolRegistry
@@ -31,6 +34,7 @@ class ChatService:
storage: JsonSessionStorage, storage: JsonSessionStorage,
client: AsyncOpenAI, client: AsyncOpenAI,
model: str, model: str,
default_system_prompt: str = "You are a helpful assistant.",
tool_registry: ToolRegistry | None = None, tool_registry: ToolRegistry | None = None,
max_tool_rounds: int = 5, max_tool_rounds: int = 5,
max_tool_calls_per_turn: int = 10, max_tool_calls_per_turn: int = 10,
@@ -40,11 +44,24 @@ class ChatService:
self.storage = storage self.storage = storage
self.client = client self.client = client
self.model = model self.model = model
self.default_system_prompt = default_system_prompt
self.tool_registry = tool_registry or ToolRegistry() self.tool_registry = tool_registry or ToolRegistry()
self.max_tool_rounds = max_tool_rounds self.max_tool_rounds = max_tool_rounds
self.max_tool_calls_per_turn = max_tool_calls_per_turn self.max_tool_calls_per_turn = max_tool_calls_per_turn
self._locks: dict[str, asyncio.Lock] = {} self._locks: dict[str, asyncio.Lock] = {}
async def create_session(
self, system_prompt: str | None = None
) -> CreateSessionResponse:
session = await self.storage.create(
system_prompt or self.default_system_prompt
)
return CreateSessionResponse(
session_id=session.session_id,
system_prompt=session.system_prompt,
created_at=session.created_at,
)
async def generate_response( async def generate_response(
self, session_id: str, content: str self, session_id: str, content: str
) -> SendMessageResponse: ) -> SendMessageResponse:
+1 -1
View File
@@ -8,7 +8,7 @@ from uuid import UUID, uuid4
from pydantic import ValidationError from pydantic import ValidationError
from models import Session from domain import Session
class SessionNotFoundError(Exception): class SessionNotFoundError(Exception):
+31 -13
View File
@@ -58,20 +58,26 @@ def test_agent_executes_tool_and_persists_full_history(tmp_path: Path) -> None:
[ [
completion( completion(
content=None, content=None,
tool_calls=[tool_call("call-1", "calculate", '{"expression":"2+2"}')], tool_calls=[
tool_call(
"call-1",
"get_current_time",
'{"timezone":"UTC"}',
)
],
), ),
completion(content="结果是 4", prompt_tokens=20, completion_tokens=3), completion(content="当前时间已查询", prompt_tokens=20, completion_tokens=3),
] ]
) )
storage = JsonSessionStorage(tmp_path) storage = JsonSessionStorage(tmp_path)
session = await storage.create("Use tools when needed.") session = await storage.create("Use tools when needed.")
service = ChatService(storage, provider, "test-model") # type: ignore[arg-type] service = ChatService(storage, provider, "test-model") # type: ignore[arg-type]
response = await service.generate_response(session.session_id, "2+2 等于多少") response = await service.generate_response(session.session_id, "现在几点")
assert response.message.content == "结果是 4" assert response.message.content == "当前时间已查询"
assert "tool_calls" not in response.message.model_dump() assert "tool_calls" not in response.message.model_dump()
assert response.tools_use == ["calculate"] assert response.tools_use == ["get_current_time"]
assert response.usage is not None assert response.usage is not None
assert response.usage.model_dump() == { assert response.usage.model_dump() == {
"prompt_tokens": 30, "prompt_tokens": 30,
@@ -86,7 +92,7 @@ def test_agent_executes_tool_and_persists_full_history(tmp_path: Path) -> None:
"assistant", "assistant",
"tool", "tool",
] ]
assert json.loads(second_messages[-1]["content"])["result"] == 4 assert json.loads(second_messages[-1]["content"])["timezone"] == "UTC"
saved = await storage.read(session.session_id) saved = await storage.read(session.session_id)
assert [message.role for message in saved.messages] == [ assert [message.role for message in saved.messages] == [
@@ -108,11 +114,15 @@ def test_agent_executes_multiple_tools_in_order(tmp_path: Path) -> None:
completion( completion(
content=None, content=None,
tool_calls=[ tool_calls=[
tool_call("call-1", "calculate", '{"expression":"6*7"}'), tool_call(
"call-1",
"get_current_time",
'{"timezone":"UTC"}',
),
tool_call( tool_call(
"call-2", "call-2",
"get_current_time", "get_current_time",
'{"timezone":"UTC"}', '{"timezone":"Asia/Shanghai"}',
), ),
], ],
), ),
@@ -130,7 +140,7 @@ def test_agent_executes_multiple_tools_in_order(tmp_path: Path) -> None:
"call-1", "call-1",
"call-2", "call-2",
] ]
assert response.tools_use == ["calculate", "get_current_time"] assert response.tools_use == ["get_current_time"]
asyncio.run(scenario()) asyncio.run(scenario())
@@ -141,7 +151,7 @@ def test_tool_error_is_returned_to_model(tmp_path: Path) -> None:
[ [
completion( completion(
content=None, content=None,
tool_calls=[tool_call("bad-1", "calculate", "not-json")], tool_calls=[tool_call("bad-1", "get_current_time", "not-json")],
), ),
completion(content="参数无效。"), completion(content="参数无效。"),
] ]
@@ -164,13 +174,21 @@ def test_tool_error_is_returned_to_model(tmp_path: Path) -> None:
@pytest.mark.parametrize( @pytest.mark.parametrize(
("max_rounds", "max_calls", "calls"), ("max_rounds", "max_calls", "calls"),
[ [
(1, 10, [tool_call("call-1", "calculate", '{"expression":"1+1"}')]), (
1,
10,
[tool_call("call-1", "get_current_time", '{"timezone":"UTC"}')],
),
( (
5, 5,
1, 1,
[ [
tool_call("call-1", "calculate", '{"expression":"1+1"}'), tool_call("call-1", "get_current_time", '{"timezone":"UTC"}'),
tool_call("call-2", "calculate", '{"expression":"2+2"}'), tool_call(
"call-2",
"get_current_time",
'{"timezone":"Asia/Shanghai"}',
),
], ],
), ),
], ],
+12 -21
View File
@@ -2,7 +2,7 @@ import asyncio
import json import json
import time import time
from tools import CalculatorArguments, ToolRegistry, ToolSpec from tools import CurrentTimeArguments, ToolRegistry, ToolSpec
def test_builtin_tool_definitions() -> None: def test_builtin_tool_definitions() -> None:
@@ -11,37 +11,28 @@ def test_builtin_tool_definitions() -> None:
definitions = registry.definitions() definitions = registry.definitions()
assert [item["function"]["name"] for item in definitions] == [ # type: ignore[index] assert [item["function"]["name"] for item in definitions] == [ # type: ignore[index]
"get_current_time", "get_current_time"
"calculate",
] ]
def test_calculator_and_time_tools() -> None: def test_current_time_tool() -> None:
async def scenario() -> None: async def scenario() -> None:
registry = ToolRegistry() registry = ToolRegistry()
calculation = await registry.execute(
"calculate", '{"expression":"(2 + 3) * 4"}'
)
current_time = await registry.execute( current_time = await registry.execute(
"get_current_time", '{"timezone":"Asia/Shanghai"}' "get_current_time", '{"timezone":"Asia/Shanghai"}'
) )
assert calculation.is_error is False
assert json.loads(calculation.content)["result"] == 20
assert current_time.is_error is False assert current_time.is_error is False
assert json.loads(current_time.content)["timezone"] == "Asia/Shanghai" assert json.loads(current_time.content)["timezone"] == "Asia/Shanghai"
asyncio.run(scenario()) asyncio.run(scenario())
def test_tool_argument_and_expression_errors_are_safe() -> None: def test_tool_argument_and_execution_errors_are_safe() -> None:
async def scenario() -> None: async def scenario() -> None:
registry = ToolRegistry() registry = ToolRegistry()
invalid_json = await registry.execute("calculate", "not-json") invalid_json = await registry.execute("get_current_time", "not-json")
unknown = await registry.execute("missing", "{}") unknown = await registry.execute("missing", "{}")
unsafe = await registry.execute(
"calculate", '{"expression":"__import__(\\"os\\").system(\\"id\\")"}'
)
invalid_timezone = await registry.execute( invalid_timezone = await registry.execute(
"get_current_time", '{"timezone":"Not/A_Real_Zone"}' "get_current_time", '{"timezone":"Not/A_Real_Zone"}'
) )
@@ -50,8 +41,6 @@ def test_tool_argument_and_expression_errors_are_safe() -> None:
assert "invalid_arguments" in invalid_json.content assert "invalid_arguments" in invalid_json.content
assert unknown.is_error is True assert unknown.is_error is True
assert "unknown_tool" in unknown.content assert "unknown_tool" in unknown.content
assert unsafe.is_error is True
assert "execution_failed" in unsafe.content
assert invalid_timezone.is_error is True assert invalid_timezone.is_error is True
assert "execution_failed" in invalid_timezone.content assert "execution_failed" in invalid_timezone.content
@@ -61,18 +50,20 @@ def test_tool_argument_and_expression_errors_are_safe() -> None:
def test_tool_timeout_returns_error() -> None: def test_tool_timeout_returns_error() -> None:
def slow_handler(arguments): def slow_handler(arguments):
time.sleep(0.05) time.sleep(0.05)
return {"result": 1} return {"timezone": "UTC"}
async def scenario() -> None: async def scenario() -> None:
registry = ToolRegistry(timeout_seconds=0.01) registry = ToolRegistry(timeout_seconds=0.01)
registry._specs["calculate"] = ToolSpec( registry._specs["get_current_time"] = ToolSpec(
name="calculate", name="get_current_time",
description="slow", description="slow",
arguments_model=CalculatorArguments, arguments_model=CurrentTimeArguments,
handler=slow_handler, handler=slow_handler,
) )
result = await registry.execute("calculate", '{"expression":"1+1"}') result = await registry.execute(
"get_current_time", '{"timezone":"UTC"}'
)
assert result.is_error is True assert result.is_error is True
assert "timeout" in result.content assert "timeout" in result.content
+1 -68
View File
@@ -1,12 +1,8 @@
import ast
import asyncio import asyncio
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
from datetime import datetime from datetime import datetime
import json import json
import math
import operator
from typing import Any
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from pydantic import BaseModel, ConfigDict, Field, ValidationError from pydantic import BaseModel, ConfigDict, Field, ValidationError
@@ -18,12 +14,6 @@ class CurrentTimeArguments(BaseModel):
timezone: str = Field(default="UTC", description="IANA timezone, e.g. Asia/Shanghai") timezone: str = Field(default="UTC", description="IANA timezone, e.g. Asia/Shanghai")
class CalculatorArguments(BaseModel):
model_config = ConfigDict(extra="forbid")
expression: str = Field(description="Arithmetic expression to evaluate")
@dataclass(frozen=True, slots=True) @dataclass(frozen=True, slots=True)
class ToolSpec: class ToolSpec:
name: str name: str
@@ -59,13 +49,7 @@ class ToolRegistry:
description="Get the current date and time in an IANA timezone.", description="Get the current date and time in an IANA timezone.",
arguments_model=CurrentTimeArguments, arguments_model=CurrentTimeArguments,
handler=_get_current_time, handler=_get_current_time,
), )
ToolSpec(
name="calculate",
description="Safely evaluate a basic arithmetic expression.",
arguments_model=CalculatorArguments,
handler=_calculate,
),
] ]
self._specs = {spec.name: spec for spec in specs} self._specs = {spec.name: spec for spec in specs}
@@ -118,54 +102,3 @@ def _get_current_time(arguments: BaseModel) -> dict[str, object]:
raise ValueError("unknown timezone") from exc raise ValueError("unknown timezone") from exc
now = datetime.now(timezone) now = datetime.now(timezone)
return {"timezone": arguments.timezone, "datetime": now.isoformat()} return {"timezone": arguments.timezone, "datetime": now.isoformat()}
_BINARY_OPERATORS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
ast.Add: operator.add,
ast.Sub: operator.sub,
ast.Mult: operator.mul,
ast.Div: operator.truediv,
ast.FloorDiv: operator.floordiv,
ast.Mod: operator.mod,
ast.Pow: operator.pow,
}
_UNARY_OPERATORS: dict[type[ast.unaryop], Callable[[Any], Any]] = {
ast.UAdd: operator.pos,
ast.USub: operator.neg,
}
_MAX_EXPRESSION_LENGTH = 200
_MAX_ABSOLUTE_RESULT = 1e100
_MAX_EXPONENT = 100
def _calculate(arguments: BaseModel) -> dict[str, object]:
assert isinstance(arguments, CalculatorArguments)
expression = arguments.expression.strip()
if not expression or len(expression) > _MAX_EXPRESSION_LENGTH:
raise ValueError("expression is empty or too long")
tree = ast.parse(expression, mode="eval")
result = _evaluate_node(tree.body)
if isinstance(result, float) and not math.isfinite(result):
raise ValueError("result is not finite")
if abs(result) > _MAX_ABSOLUTE_RESULT:
raise ValueError("result is too large")
return {"expression": expression, "result": result}
def _evaluate_node(node: ast.AST) -> int | float:
if isinstance(node, ast.Constant):
if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
raise ValueError("only numeric constants are allowed")
return node.value
if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPERATORS:
return _UNARY_OPERATORS[type(node.op)](_evaluate_node(node.operand))
if isinstance(node, ast.BinOp) and type(node.op) in _BINARY_OPERATORS:
left = _evaluate_node(node.left)
right = _evaluate_node(node.right)
if isinstance(node.op, ast.Pow) and abs(right) > _MAX_EXPONENT:
raise ValueError("exponent is too large")
result = _BINARY_OPERATORS[type(node.op)](left, right)
if abs(result) > _MAX_ABSOLUTE_RESULT:
raise ValueError("intermediate result is too large")
return result
raise ValueError("unsupported expression")