mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
Redesign Python AI engine (#5991)
# Description of Changes Redesign the Python AI engine to be properly agentic and make use of `pydantic-ai` instead of `langchain` for correctness and ergonomics. This should be a good foundation for us to build our AI engine on going forwards.
This commit is contained in:
James Brunton
@@ -0,0 +1,7 @@
|
||||
"""Stirling AI engine package."""
|
||||
|
||||
from .api.app import app
|
||||
|
||||
__all__ = [
|
||||
"app",
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Agent modules for Stirling AI reasoning flows."""
|
||||
|
||||
from .execution import ExecutionPlanningAgent
|
||||
from .orchestrator import OrchestratorAgent
|
||||
from .pdf_edit import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
|
||||
from .pdf_questions import PdfQuestionAgent
|
||||
from .user_spec import UserSpecAgent
|
||||
|
||||
__all__ = [
|
||||
"ExecutionPlanningAgent",
|
||||
"OrchestratorAgent",
|
||||
"PdfEditAgent",
|
||||
"PdfEditParameterSelector",
|
||||
"PdfEditPlanSelection",
|
||||
"PdfQuestionAgent",
|
||||
"UserSpecAgent",
|
||||
]
|
||||
@@ -0,0 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.contracts import AgentExecutionRequest, CannotContinueExecutionAction, NextExecutionAction
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class ExecutionPlanningAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
|
||||
async def next_action(self, request: AgentExecutionRequest) -> NextExecutionAction:
|
||||
return CannotContinueExecutionAction(
|
||||
reason=f"Execution planning is not implemented yet for step {request.current_step_index}."
|
||||
)
|
||||
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import ToolOutput
|
||||
from pydantic_ai.tools import RunContext
|
||||
|
||||
from stirling.agents.pdf_edit import PdfEditAgent
|
||||
from stirling.agents.pdf_questions import PdfQuestionAgent
|
||||
from stirling.agents.user_spec import UserSpecAgent
|
||||
from stirling.contracts import (
|
||||
AgentDraftRequest,
|
||||
AgentDraftWorkflowResponse,
|
||||
OrchestratorRequest,
|
||||
OrchestratorResponse,
|
||||
PdfEditRequest,
|
||||
PdfEditResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
UnsupportedCapabilityResponse,
|
||||
)
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OrchestratorDeps:
|
||||
runtime: AppRuntime
|
||||
request: OrchestratorRequest
|
||||
|
||||
|
||||
class OrchestratorAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.agent = Agent(
|
||||
model=runtime.fast_model,
|
||||
output_type=[
|
||||
ToolOutput(
|
||||
self.delegate_pdf_edit,
|
||||
name="delegate_pdf_edit",
|
||||
description="Delegate requests for PDF modifications and return the PDF edit result.",
|
||||
),
|
||||
ToolOutput(
|
||||
self.delegate_pdf_question,
|
||||
name="delegate_pdf_question",
|
||||
description="Delegate questions about PDF contents and return the PDF question result.",
|
||||
),
|
||||
ToolOutput(
|
||||
self.delegate_user_spec,
|
||||
name="delegate_user_spec",
|
||||
description="Delegate requests to create or revise a user agent spec and return the draft result.",
|
||||
),
|
||||
ToolOutput(
|
||||
self.unsupported_capability,
|
||||
name="unsupported_capability",
|
||||
description="Return this when none of the delegate outputs fit the request.",
|
||||
),
|
||||
],
|
||||
deps_type=OrchestratorDeps,
|
||||
system_prompt=(
|
||||
"You are the top-level orchestrator. "
|
||||
"Choose exactly one output function that best handles the request. "
|
||||
"Use delegate_pdf_edit for requested PDF modifications. "
|
||||
"Use delegate_pdf_question for questions about the contents of a PDF. "
|
||||
"Use delegate_user_spec for requests to create or define an agent spec. "
|
||||
"Use unsupported_capability only when none of the other outputs fit."
|
||||
),
|
||||
model_settings=runtime.fast_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
|
||||
result = await self.agent.run(
|
||||
request.user_message,
|
||||
deps=OrchestratorDeps(runtime=self.runtime, request=request),
|
||||
)
|
||||
return result.output
|
||||
|
||||
async def delegate_pdf_edit(self, ctx: RunContext[OrchestratorDeps]) -> PdfEditResponse:
|
||||
request = ctx.deps.request
|
||||
return await PdfEditAgent(ctx.deps.runtime).handle(
|
||||
PdfEditRequest(user_message=request.user_message, conversation_id=request.conversation_id)
|
||||
)
|
||||
|
||||
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
|
||||
request = ctx.deps.request
|
||||
return await PdfQuestionAgent(ctx.deps.runtime).handle(
|
||||
PdfQuestionRequest(question=request.user_message, conversation_id=request.conversation_id)
|
||||
)
|
||||
|
||||
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
|
||||
request = ctx.deps.request
|
||||
return await UserSpecAgent(ctx.deps.runtime).draft(AgentDraftRequest(user_message=request.user_message))
|
||||
|
||||
async def unsupported_capability(
|
||||
self,
|
||||
ctx: RunContext[OrchestratorDeps],
|
||||
capability: str,
|
||||
message: str,
|
||||
) -> UnsupportedCapabilityResponse:
|
||||
return UnsupportedCapabilityResponse(capability=capability, message=message)
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.contracts import (
|
||||
EditCannotDoResponse,
|
||||
EditClarificationRequest,
|
||||
EditPlanResponse,
|
||||
PdfEditRequest,
|
||||
PdfEditResponse,
|
||||
ToolOperationStep,
|
||||
)
|
||||
from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class PdfEditPlanSelection(ApiModel):
|
||||
outcome: Literal["plan"] = "plan"
|
||||
operations: list[OperationId] = Field(min_length=1)
|
||||
summary: str
|
||||
rationale: str | None = None
|
||||
|
||||
|
||||
class PdfEditParameterSelector:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
system_prompt=(
|
||||
"Generate only the parameter object for the selected PDF operation. "
|
||||
"Use reasonable defaults when the request does not specify optional details. "
|
||||
"Only fill fields that belong to the selected operation's parameter model."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def select(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
operation_plan: list[OperationId],
|
||||
operation_index: int,
|
||||
generated_steps: list[ToolOperationStep],
|
||||
) -> ParamToolModel:
|
||||
operation_id = operation_plan[operation_index]
|
||||
parameter_model = OPERATIONS[operation_id]
|
||||
parameter_result = await self.agent.run(
|
||||
self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps),
|
||||
output_type=NativeOutput(parameter_model),
|
||||
instructions=(
|
||||
f"Generate only the parameters for the PDF operation `{operation_id.value}`. "
|
||||
"Do not include fields from any other operation."
|
||||
),
|
||||
)
|
||||
return parameter_result.output
|
||||
|
||||
def _build_parameter_prompt(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
operation_plan: list[OperationId],
|
||||
operation_index: int,
|
||||
generated_steps: list[ToolOperationStep],
|
||||
) -> str:
|
||||
operation_id = operation_plan[operation_index]
|
||||
operation_list = ", ".join(operation.value for operation in operation_plan)
|
||||
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
|
||||
generated_steps_text = (
|
||||
"\n".join(
|
||||
f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps)
|
||||
)
|
||||
if generated_steps
|
||||
else "None"
|
||||
)
|
||||
return (
|
||||
f"User request: {request.user_message}\n"
|
||||
f"Files: {file_names}\n"
|
||||
f"Operation plan: {operation_list}\n"
|
||||
f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
|
||||
f"Selected operation: {operation_id.value}\n"
|
||||
f"Already generated steps:\n{generated_steps_text}\n"
|
||||
"Return only the parameter object for the selected operation."
|
||||
)
|
||||
|
||||
|
||||
class PdfEditAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.supported_operations = list(OPERATIONS)
|
||||
self.parameter_selector = PdfEditParameterSelector(runtime)
|
||||
self.selection_agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
output_type=NativeOutput(
|
||||
[
|
||||
PdfEditPlanSelection,
|
||||
EditClarificationRequest,
|
||||
EditCannotDoResponse,
|
||||
]
|
||||
),
|
||||
system_prompt=(
|
||||
"Plan PDF edit requests. "
|
||||
f"Supported operations are: {self._supported_operations_prompt()}. "
|
||||
"Return an ordered list of one or more supported operations for the plan. "
|
||||
"Do not produce operation parameters in this stage. "
|
||||
"Return need_clarification when the request is genuinely ambiguous. "
|
||||
"Return cannot_do when the request is outside the supported operations. "
|
||||
"Return plan when a reasonable multi-step plan can be created. "
|
||||
"Never return partial plans."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: PdfEditRequest) -> PdfEditResponse:
|
||||
selection = await self._select_plan(request)
|
||||
if isinstance(selection, EditClarificationRequest | EditCannotDoResponse):
|
||||
return selection
|
||||
steps: list[ToolOperationStep] = []
|
||||
for operation_index, operation_id in enumerate(selection.operations):
|
||||
parameters = await self.parameter_selector.select(
|
||||
request,
|
||||
selection.operations,
|
||||
operation_index,
|
||||
steps,
|
||||
)
|
||||
steps.append(
|
||||
ToolOperationStep(
|
||||
tool=operation_id,
|
||||
parameters=parameters,
|
||||
)
|
||||
)
|
||||
return EditPlanResponse(
|
||||
summary=selection.summary,
|
||||
rationale=selection.rationale,
|
||||
steps=steps,
|
||||
)
|
||||
|
||||
async def _select_plan(
|
||||
self,
|
||||
request: PdfEditRequest,
|
||||
) -> PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse:
|
||||
selection_result = await self.selection_agent.run(self._build_selection_prompt(request))
|
||||
return selection_result.output
|
||||
|
||||
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
|
||||
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
|
||||
return (
|
||||
f"User request: {request.user_message}\n"
|
||||
f"Files: {file_names}\n"
|
||||
f"Supported operations: {self._supported_operations_prompt()}\n"
|
||||
"Plan an ordered list of supported PDF edit operations or return clarification/cannot_do."
|
||||
)
|
||||
|
||||
def _supported_operations_prompt(self) -> str:
|
||||
return ", ".join(operation_id.value for operation_id in self.supported_operations)
|
||||
@@ -0,0 +1,49 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.contracts import (
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
)
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class PdfQuestionAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
output_type=NativeOutput(
|
||||
[
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
]
|
||||
),
|
||||
system_prompt=(
|
||||
"Answer questions about a PDF using only the extracted text provided in the prompt. "
|
||||
"Do not guess or use outside knowledge. "
|
||||
"If the answer is not supported by the provided text, return not_found. "
|
||||
"When answering, include a short list of evidence snippets copied from the provided text."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||
if not request.extracted_text.strip():
|
||||
return PdfQuestionNeedTextResponse(
|
||||
reason="No extracted PDF text was provided, so the question cannot be answered yet."
|
||||
)
|
||||
return await self._run_answer_agent(request)
|
||||
|
||||
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||
result = await self.agent.run(self._build_prompt(request))
|
||||
return result.output
|
||||
|
||||
def _build_prompt(self, request: PdfQuestionRequest) -> str:
|
||||
file_name = request.file_name or "Unknown file"
|
||||
return f"File: {file_name}\nQuestion: {request.question}\nExtracted text:\n{request.extracted_text}"
|
||||
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.agents.pdf_edit import PdfEditAgent
|
||||
from stirling.contracts import (
|
||||
AgentDraft,
|
||||
AgentDraftRequest,
|
||||
AgentDraftResponse,
|
||||
AgentDraftWorkflowResponse,
|
||||
AgentRevisionRequest,
|
||||
AgentRevisionResponse,
|
||||
AgentRevisionWorkflowResponse,
|
||||
AiToolAgentStep,
|
||||
ConversationMessage,
|
||||
EditCannotDoResponse,
|
||||
EditClarificationRequest,
|
||||
EditPlanResponse,
|
||||
PdfEditRequest,
|
||||
)
|
||||
from stirling.models import ApiModel
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
class UserSpecMetadata(ApiModel):
|
||||
name: str
|
||||
description: str
|
||||
objective: str
|
||||
|
||||
|
||||
class UserSpecAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
self.runtime = runtime
|
||||
self.pdf_edit_agent = PdfEditAgent(runtime)
|
||||
self.agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
output_type=NativeOutput(UserSpecMetadata),
|
||||
system_prompt=(
|
||||
"Create or revise a saved agent draft from the provided request and edit plan. "
|
||||
"Return a concise name, description, and objective. "
|
||||
"Keep the workflow grounded and practical."
|
||||
),
|
||||
model_settings=runtime.smart_model_settings,
|
||||
)
|
||||
|
||||
async def draft(self, request: AgentDraftRequest) -> AgentDraftWorkflowResponse:
|
||||
edit_plan = await self._build_edit_plan(request.user_message)
|
||||
if not isinstance(edit_plan, EditPlanResponse):
|
||||
return edit_plan
|
||||
return AgentDraftResponse(draft=await self._run_draft_agent(request, edit_plan))
|
||||
|
||||
async def revise(self, request: AgentRevisionRequest) -> AgentRevisionWorkflowResponse:
|
||||
edit_plan = await self._build_edit_plan(
|
||||
f"Current objective: {request.current_draft.objective}\nRevision request: {request.user_message}"
|
||||
)
|
||||
if not isinstance(edit_plan, EditPlanResponse):
|
||||
return edit_plan
|
||||
return AgentRevisionResponse(draft=await self._run_revision_agent(request, edit_plan))
|
||||
|
||||
async def _run_draft_agent(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> AgentDraft:
|
||||
metadata = (await self.agent.run(self._build_draft_prompt(request, edit_plan))).output
|
||||
return AgentDraft(
|
||||
name=metadata.name,
|
||||
description=metadata.description,
|
||||
objective=metadata.objective,
|
||||
steps=[*edit_plan.steps],
|
||||
)
|
||||
|
||||
async def _run_revision_agent(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> AgentDraft:
|
||||
metadata = (await self.agent.run(self._build_revision_prompt(request, edit_plan))).output
|
||||
preserved_ai_steps = [step for step in request.current_draft.steps if isinstance(step, AiToolAgentStep)]
|
||||
return AgentDraft(
|
||||
name=metadata.name,
|
||||
description=metadata.description,
|
||||
objective=metadata.objective,
|
||||
steps=[*edit_plan.steps, *preserved_ai_steps],
|
||||
)
|
||||
|
||||
def _build_draft_prompt(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> str:
|
||||
return (
|
||||
f"User request:\n{request.user_message}\n\n"
|
||||
f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n"
|
||||
f"Edit plan summary:\n{edit_plan.summary}\n\n"
|
||||
f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n"
|
||||
f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}"
|
||||
)
|
||||
|
||||
def _build_revision_prompt(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> str:
|
||||
return (
|
||||
f"Revision request:\n{request.user_message}\n\n"
|
||||
f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n"
|
||||
f"Current draft:\n{request.current_draft.model_dump_json(indent=2)}\n\n"
|
||||
f"Edit plan summary:\n{edit_plan.summary}\n\n"
|
||||
f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n"
|
||||
f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}"
|
||||
)
|
||||
|
||||
def _format_conversation_history(self, conversation_history: list[ConversationMessage]) -> str:
|
||||
if not conversation_history:
|
||||
return "None"
|
||||
return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
|
||||
|
||||
async def _build_edit_plan(
|
||||
self,
|
||||
user_message: str,
|
||||
) -> EditPlanResponse | EditClarificationRequest | EditCannotDoResponse:
|
||||
return await self.pdf_edit_agent.handle(PdfEditRequest(user_message=user_message))
|
||||
@@ -0,0 +1,7 @@
|
||||
"""API surface for the Stirling AI service."""
|
||||
|
||||
from .app import app
|
||||
|
||||
__all__ = [
|
||||
"app",
|
||||
]
|
||||
@@ -0,0 +1,57 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import Depends, FastAPI
|
||||
|
||||
from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
|
||||
from stirling.api.routes import (
|
||||
agent_draft_router,
|
||||
execution_router,
|
||||
orchestrator_router,
|
||||
pdf_edit_router,
|
||||
pdf_question_router,
|
||||
)
|
||||
from stirling.config import AppSettings, load_settings
|
||||
from stirling.contracts import HealthResponse
|
||||
from stirling.services import build_runtime
|
||||
|
||||
|
||||
def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
|
||||
override = fast_api.dependency_overrides.get(load_settings)
|
||||
if override is not None:
|
||||
return override()
|
||||
return load_settings()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(fast_api: FastAPI):
|
||||
# Load env vars on startup so we can immediately crash if required env vars aren't set
|
||||
settings = _load_startup_settings(fast_api)
|
||||
runtime = build_runtime(settings)
|
||||
fast_api.state.settings = settings
|
||||
fast_api.state.runtime = runtime
|
||||
fast_api.state.orchestrator_agent = OrchestratorAgent(runtime)
|
||||
fast_api.state.pdf_edit_agent = PdfEditAgent(runtime)
|
||||
fast_api.state.pdf_question_agent = PdfQuestionAgent(runtime)
|
||||
fast_api.state.user_spec_agent = UserSpecAgent(runtime)
|
||||
fast_api.state.execution_planning_agent = ExecutionPlanningAgent(runtime)
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
|
||||
app.include_router(orchestrator_router)
|
||||
app.include_router(pdf_edit_router)
|
||||
app.include_router(pdf_question_router)
|
||||
app.include_router(agent_draft_router)
|
||||
app.include_router(execution_router)
|
||||
|
||||
|
||||
@app.get("/health", response_model=HealthResponse)
|
||||
async def healthcheck(settings: Annotated[AppSettings, Depends(load_settings)]) -> HealthResponse:
|
||||
return HealthResponse(
|
||||
status="ok",
|
||||
smart_model=settings.smart_model_name,
|
||||
fast_model=settings.fast_model_name,
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
def get_runtime(request: Request) -> AppRuntime:
|
||||
return request.app.state.runtime
|
||||
|
||||
|
||||
def get_orchestrator_agent(request: Request) -> OrchestratorAgent:
|
||||
return request.app.state.orchestrator_agent
|
||||
|
||||
|
||||
def get_pdf_edit_agent(request: Request) -> PdfEditAgent:
|
||||
return request.app.state.pdf_edit_agent
|
||||
|
||||
|
||||
def get_pdf_question_agent(request: Request) -> PdfQuestionAgent:
|
||||
return request.app.state.pdf_question_agent
|
||||
|
||||
|
||||
def get_user_spec_agent(request: Request) -> UserSpecAgent:
|
||||
return request.app.state.user_spec_agent
|
||||
|
||||
|
||||
def get_execution_planning_agent(request: Request) -> ExecutionPlanningAgent:
|
||||
return request.app.state.execution_planning_agent
|
||||
@@ -0,0 +1,13 @@
|
||||
from .agent_drafts import router as agent_draft_router
|
||||
from .execution import router as execution_router
|
||||
from .orchestrator import router as orchestrator_router
|
||||
from .pdf_edit import router as pdf_edit_router
|
||||
from .pdf_questions import router as pdf_question_router
|
||||
|
||||
__all__ = [
|
||||
"agent_draft_router",
|
||||
"execution_router",
|
||||
"orchestrator_router",
|
||||
"pdf_edit_router",
|
||||
"pdf_question_router",
|
||||
]
|
||||
@@ -0,0 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.agents import UserSpecAgent
|
||||
from stirling.api.dependencies import get_user_spec_agent
|
||||
from stirling.contracts import (
|
||||
AgentDraftRequest,
|
||||
AgentDraftWorkflowResponse,
|
||||
AgentRevisionRequest,
|
||||
AgentRevisionWorkflowResponse,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/api/v1/agents", tags=["agents"])
|
||||
|
||||
|
||||
@router.post("/draft", response_model=AgentDraftWorkflowResponse)
|
||||
async def draft_agent(
|
||||
request: AgentDraftRequest,
|
||||
agent: Annotated[UserSpecAgent, Depends(get_user_spec_agent)],
|
||||
) -> AgentDraftWorkflowResponse:
|
||||
return await agent.draft(request)
|
||||
|
||||
|
||||
@router.post("/revise", response_model=AgentRevisionWorkflowResponse)
|
||||
async def revise_agent(
|
||||
request: AgentRevisionRequest,
|
||||
agent: Annotated[UserSpecAgent, Depends(get_user_spec_agent)],
|
||||
) -> AgentRevisionWorkflowResponse:
|
||||
return await agent.revise(request)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.agents import ExecutionPlanningAgent
|
||||
from stirling.api.dependencies import get_execution_planning_agent
|
||||
from stirling.contracts import AgentExecutionRequest, NextExecutionAction
|
||||
|
||||
router = APIRouter(prefix="/api/v1/agents", tags=["agents"])
|
||||
|
||||
|
||||
@router.post("/next-action", response_model=NextExecutionAction)
|
||||
async def next_action(
|
||||
request: AgentExecutionRequest,
|
||||
agent: Annotated[ExecutionPlanningAgent, Depends(get_execution_planning_agent)],
|
||||
) -> NextExecutionAction:
|
||||
return await agent.next_action(request)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.agents import OrchestratorAgent
|
||||
from stirling.api.dependencies import get_orchestrator_agent
|
||||
from stirling.contracts import OrchestratorRequest, OrchestratorResponse
|
||||
|
||||
router = APIRouter(prefix="/api/v1/orchestrator", tags=["orchestrator"])
|
||||
|
||||
|
||||
@router.post("", response_model=OrchestratorResponse)
|
||||
async def orchestrate(
|
||||
request: OrchestratorRequest,
|
||||
agent: Annotated[OrchestratorAgent, Depends(get_orchestrator_agent)],
|
||||
) -> OrchestratorResponse:
|
||||
return await agent.handle(request)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.agents import PdfEditAgent
|
||||
from stirling.api.dependencies import get_pdf_edit_agent
|
||||
from stirling.contracts import PdfEditRequest, PdfEditResponse
|
||||
|
||||
router = APIRouter(prefix="/api/v1/pdf/edit", tags=["pdf-edit"])
|
||||
|
||||
|
||||
@router.post("", response_model=PdfEditResponse)
|
||||
async def pdf_edit(
|
||||
request: PdfEditRequest,
|
||||
agent: Annotated[PdfEditAgent, Depends(get_pdf_edit_agent)],
|
||||
) -> PdfEditResponse:
|
||||
return await agent.handle(request)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.agents import PdfQuestionAgent
|
||||
from stirling.api.dependencies import get_pdf_question_agent
|
||||
from stirling.contracts import PdfQuestionRequest, PdfQuestionResponse
|
||||
|
||||
router = APIRouter(prefix="/api/v1/pdf/questions", tags=["pdf-questions"])
|
||||
|
||||
|
||||
@router.post("", response_model=PdfQuestionResponse)
|
||||
async def pdf_questions(
|
||||
request: PdfQuestionRequest,
|
||||
agent: Annotated[PdfQuestionAgent, Depends(get_pdf_question_agent)],
|
||||
) -> PdfQuestionResponse:
|
||||
return await agent.handle(request)
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Configuration models and loaders for the Stirling AI service."""
|
||||
|
||||
from .settings import AppSettings, load_settings
|
||||
|
||||
__all__ = [
|
||||
"AppSettings",
|
||||
"load_settings",
|
||||
]
|
||||
@@ -0,0 +1,26 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
ENGINE_ROOT = Path(__file__).resolve().parents[3]
|
||||
ENV_FILE = ENGINE_ROOT / ".env"
|
||||
|
||||
|
||||
class AppSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(env_file=ENV_FILE, extra="ignore", populate_by_name=True)
|
||||
|
||||
smart_model_name: str = Field(validation_alias="STIRLING_SMART_MODEL")
|
||||
fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL")
|
||||
smart_model_max_tokens: int = Field(validation_alias="STIRLING_SMART_MODEL_MAX_TOKENS")
|
||||
fast_model_max_tokens: int = Field(validation_alias="STIRLING_FAST_MODEL_MAX_TOKENS")
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def load_settings() -> AppSettings:
|
||||
load_dotenv(ENV_FILE)
|
||||
return AppSettings.model_validate({})
|
||||
@@ -0,0 +1,74 @@
|
||||
from .agent_drafts import (
|
||||
AgentDraft,
|
||||
AgentDraftRequest,
|
||||
AgentDraftResponse,
|
||||
AgentDraftWorkflowResponse,
|
||||
AgentRevisionRequest,
|
||||
AgentRevisionResponse,
|
||||
AgentRevisionWorkflowResponse,
|
||||
)
|
||||
from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
|
||||
from .common import ConversationMessage, PdfTextSelection, ToolOperationStep
|
||||
from .execution import (
|
||||
AgentExecutionRequest,
|
||||
CannotContinueExecutionAction,
|
||||
CompletedExecutionAction,
|
||||
ExecutionContext,
|
||||
ExecutionStepResult,
|
||||
NextExecutionAction,
|
||||
ToolCallExecutionAction,
|
||||
)
|
||||
from .health import HealthResponse
|
||||
from .orchestrator import OrchestratorRequest, OrchestratorResponse, SupportedCapability, UnsupportedCapabilityResponse
|
||||
from .pdf_edit import (
|
||||
EditCannotDoResponse,
|
||||
EditClarificationRequest,
|
||||
EditPlanResponse,
|
||||
PdfEditRequest,
|
||||
PdfEditResponse,
|
||||
)
|
||||
from .pdf_questions import (
|
||||
PdfQuestionAnswerResponse,
|
||||
PdfQuestionNeedTextResponse,
|
||||
PdfQuestionNotFoundResponse,
|
||||
PdfQuestionRequest,
|
||||
PdfQuestionResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AgentDraft",
|
||||
"AgentDraftRequest",
|
||||
"AgentDraftResponse",
|
||||
"AgentDraftWorkflowResponse",
|
||||
"AgentExecutionRequest",
|
||||
"AgentRevisionRequest",
|
||||
"AgentRevisionResponse",
|
||||
"AgentRevisionWorkflowResponse",
|
||||
"AgentSpec",
|
||||
"AgentSpecStep",
|
||||
"AiToolAgentStep",
|
||||
"CannotContinueExecutionAction",
|
||||
"ConversationMessage",
|
||||
"CompletedExecutionAction",
|
||||
"EditCannotDoResponse",
|
||||
"EditClarificationRequest",
|
||||
"EditPlanResponse",
|
||||
"ExecutionContext",
|
||||
"ExecutionStepResult",
|
||||
"HealthResponse",
|
||||
"NextExecutionAction",
|
||||
"OrchestratorRequest",
|
||||
"OrchestratorResponse",
|
||||
"PdfEditRequest",
|
||||
"PdfEditResponse",
|
||||
"PdfQuestionAnswerResponse",
|
||||
"PdfQuestionNotFoundResponse",
|
||||
"PdfQuestionNeedTextResponse",
|
||||
"PdfQuestionRequest",
|
||||
"PdfQuestionResponse",
|
||||
"PdfTextSelection",
|
||||
"SupportedCapability",
|
||||
"ToolOperationStep",
|
||||
"ToolCallExecutionAction",
|
||||
"UnsupportedCapabilityResponse",
|
||||
]
|
||||
@@ -0,0 +1,57 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .agent_specs import AgentSpecStep
|
||||
from .common import ConversationMessage
|
||||
from .pdf_edit import EditCannotDoResponse, EditClarificationRequest
|
||||
|
||||
|
||||
class AgentDraftStep(ApiModel):
|
||||
kind: Literal["tool", "ai_tool"]
|
||||
title: str
|
||||
description: str
|
||||
|
||||
|
||||
class AgentDraft(ApiModel):
|
||||
name: str
|
||||
description: str
|
||||
objective: str
|
||||
steps: list[AgentSpecStep] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentDraftRequest(ApiModel):
|
||||
user_message: str
|
||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||
|
||||
|
||||
class AgentDraftResponse(ApiModel):
|
||||
outcome: Literal["draft"] = "draft"
|
||||
draft: AgentDraft
|
||||
|
||||
|
||||
class AgentRevisionRequest(ApiModel):
|
||||
user_message: str
|
||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||
current_draft: AgentDraft
|
||||
|
||||
|
||||
class AgentRevisionResponse(ApiModel):
|
||||
outcome: Literal["draft"] = "draft"
|
||||
draft: AgentDraft
|
||||
|
||||
|
||||
AgentDraftWorkflowResponse = Annotated[
|
||||
AgentDraftResponse | EditClarificationRequest | EditCannotDoResponse,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
|
||||
|
||||
AgentRevisionWorkflowResponse = Annotated[
|
||||
AgentRevisionResponse | EditClarificationRequest | EditCannotDoResponse,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel, OperationId
|
||||
|
||||
from .common import ToolOperationStep
|
||||
|
||||
|
||||
class AiToolAgentStep(ApiModel):
|
||||
kind: Literal["ai_tool"] = "ai_tool"
|
||||
title: str
|
||||
description: str
|
||||
tool: OperationId
|
||||
instruction: str
|
||||
|
||||
|
||||
AgentSpecStep = Annotated[ToolOperationStep | AiToolAgentStep, Field(discriminator="kind")]
|
||||
|
||||
|
||||
class AgentSpec(ApiModel):
|
||||
name: str
|
||||
description: str
|
||||
objective: str
|
||||
steps: list[AgentSpecStep] = Field(default_factory=list)
|
||||
@@ -0,0 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import model_validator
|
||||
|
||||
from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
|
||||
|
||||
|
||||
class ConversationMessage(ApiModel):
|
||||
role: str
|
||||
content: str
|
||||
|
||||
|
||||
class PdfTextSelection(ApiModel):
|
||||
page_number: int | None = None
|
||||
text: str
|
||||
|
||||
|
||||
class ToolOperationStep(ApiModel):
|
||||
kind: Literal["tool"] = "tool"
|
||||
tool: OperationId
|
||||
parameters: ParamToolModel
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_tool_parameter_pairing(self) -> ToolOperationStep:
|
||||
expected_type = OPERATIONS[self.tool]
|
||||
if not isinstance(self.parameters, expected_type):
|
||||
actual_type = type(self.parameters).__name__
|
||||
expected_type_name = expected_type.__name__
|
||||
raise ValueError(f"Parameters for tool {self.tool.value} must be {expected_type_name}, got {actual_type}.")
|
||||
return self
|
||||
@@ -0,0 +1,53 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel, OperationId, ParamToolModel
|
||||
|
||||
from .agent_specs import AgentSpec
|
||||
|
||||
|
||||
class ExecutionStepResult(ApiModel):
|
||||
step_index: int
|
||||
tool: OperationId | None = None
|
||||
success: bool
|
||||
output_summary: str | None = None
|
||||
output_data: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ExecutionContext(ApiModel):
|
||||
trigger_type: str | None = None
|
||||
input_files: list[str] = Field(default_factory=list)
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class AgentExecutionRequest(ApiModel):
|
||||
agent_spec: AgentSpec
|
||||
current_step_index: int
|
||||
execution_context: ExecutionContext
|
||||
previous_step_results: list[ExecutionStepResult] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ToolCallExecutionAction(ApiModel):
|
||||
outcome: Literal["tool_call"] = "tool_call"
|
||||
tool: OperationId
|
||||
parameters: ParamToolModel
|
||||
rationale: str | None = None
|
||||
|
||||
|
||||
class CompletedExecutionAction(ApiModel):
|
||||
outcome: Literal["completed"] = "completed"
|
||||
summary: str
|
||||
|
||||
|
||||
class CannotContinueExecutionAction(ApiModel):
|
||||
outcome: Literal["cannot_continue"] = "cannot_continue"
|
||||
reason: str
|
||||
|
||||
|
||||
NextExecutionAction = Annotated[
|
||||
ToolCallExecutionAction | CompletedExecutionAction | CannotContinueExecutionAction,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
@@ -0,0 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
|
||||
class HealthResponse(ApiModel):
|
||||
status: str
|
||||
smart_model: str
|
||||
fast_model: str
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .agent_drafts import AgentDraftResponse
|
||||
from .execution import NextExecutionAction
|
||||
from .pdf_edit import PdfEditResponse
|
||||
from .pdf_questions import PdfQuestionResponse
|
||||
|
||||
|
||||
class SupportedCapability(StrEnum):
|
||||
ORCHESTRATE = "orchestrate"
|
||||
PDF_EDIT = "pdf_edit"
|
||||
PDF_QUESTION = "pdf_question"
|
||||
AGENT_DRAFT = "agent_draft"
|
||||
AGENT_REVISE = "agent_revise"
|
||||
AGENT_NEXT_ACTION = "agent_next_action"
|
||||
|
||||
|
||||
class OrchestratorRequest(ApiModel):
|
||||
user_message: str
|
||||
conversation_id: str | None = None
|
||||
|
||||
|
||||
class UnsupportedCapabilityResponse(ApiModel):
|
||||
outcome: Literal["unsupported_capability"] = "unsupported_capability"
|
||||
capability: str
|
||||
message: str
|
||||
|
||||
|
||||
OrchestratorResponse = Annotated[
|
||||
PdfEditResponse | PdfQuestionResponse | AgentDraftResponse | NextExecutionAction | UnsupportedCapabilityResponse,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .common import ToolOperationStep
|
||||
|
||||
|
||||
class PdfEditRequest(ApiModel):
|
||||
user_message: str
|
||||
conversation_id: str | None = None
|
||||
file_names: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class EditPlanResponse(ApiModel):
|
||||
outcome: Literal["plan"] = "plan"
|
||||
summary: str
|
||||
rationale: str | None = None
|
||||
steps: list[ToolOperationStep]
|
||||
|
||||
|
||||
class EditClarificationRequest(ApiModel):
|
||||
outcome: Literal["need_clarification"] = "need_clarification"
|
||||
question: str
|
||||
reason: str
|
||||
|
||||
|
||||
class EditCannotDoResponse(ApiModel):
|
||||
outcome: Literal["cannot_do"] = "cannot_do"
|
||||
reason: str
|
||||
|
||||
|
||||
PdfEditResponse = Annotated[
|
||||
EditPlanResponse | EditClarificationRequest | EditCannotDoResponse,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
|
||||
class PdfQuestionRequest(ApiModel):
|
||||
question: str
|
||||
conversation_id: str | None = None
|
||||
extracted_text: str = ""
|
||||
file_name: str | None = None
|
||||
|
||||
|
||||
class PdfQuestionAnswerResponse(ApiModel):
|
||||
outcome: Literal["answer"] = "answer"
|
||||
answer: str
|
||||
evidence: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class PdfQuestionNeedTextResponse(ApiModel):
|
||||
outcome: Literal["need_text"] = "need_text"
|
||||
reason: str
|
||||
|
||||
|
||||
class PdfQuestionNotFoundResponse(ApiModel):
|
||||
outcome: Literal["not_found"] = "not_found"
|
||||
reason: str
|
||||
|
||||
|
||||
PdfQuestionResponse = Annotated[
|
||||
PdfQuestionAnswerResponse | PdfQuestionNeedTextResponse | PdfQuestionNotFoundResponse,
|
||||
Field(discriminator="outcome"),
|
||||
]
|
||||
@@ -0,0 +1,11 @@
|
||||
from . import tool_models
|
||||
from .base import ApiModel
|
||||
from .tool_models import OPERATIONS, OperationId, ParamToolModel
|
||||
|
||||
__all__ = [
|
||||
"ApiModel",
|
||||
"OPERATIONS",
|
||||
"OperationId",
|
||||
"ParamToolModel",
|
||||
"tool_models",
|
||||
]
|
||||
@@ -0,0 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from pydantic.alias_generators import to_camel
|
||||
|
||||
|
||||
class ApiModel(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
alias_generator=to_camel,
|
||||
extra="forbid",
|
||||
validate_by_name=True,
|
||||
validate_by_alias=True,
|
||||
)
|
||||
@@ -0,0 +1,439 @@
|
||||
# AUTO-GENERATED FILE. DO NOT EDIT.
|
||||
# Generated by scripts/generate_tool_models.py from frontend TypeScript sources.
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Any, Literal
|
||||
|
||||
from stirling.models.base import ApiModel
|
||||
|
||||
|
||||
class AddAttachmentsParams(ApiModel):
|
||||
attachments: list[dict[str, Any]] = []
|
||||
|
||||
|
||||
class AdjustContrastParams(ApiModel):
|
||||
blue: float = 100
|
||||
brightness: float = 100
|
||||
contrast: float = 100
|
||||
green: float = 100
|
||||
red: float = 100
|
||||
saturation: float = 100
|
||||
|
||||
|
||||
class AutoRenameParams(ApiModel):
|
||||
use_first_text_as_fallback: bool = False
|
||||
|
||||
|
||||
class AutomateParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class BookletImpositionParams(ApiModel):
|
||||
add_border: bool = False
|
||||
add_gutter: bool = False
|
||||
double_sided: bool = True
|
||||
duplex_pass: Literal["BOTH", "FIRST", "SECOND"] = "BOTH"
|
||||
flip_on_short_edge: bool = False
|
||||
gutter_size: float = 12
|
||||
pages_per_sheet: float = 2
|
||||
spine_location: Literal["LEFT", "RIGHT"] = "LEFT"
|
||||
|
||||
|
||||
class CertSignParams(ApiModel):
|
||||
cert_file: dict[str, Any] | None = None
|
||||
cert_type: Literal["", "PEM", "PKCS12", "PFX", "JKS"] = ""
|
||||
jks_file: dict[str, Any] | None = None
|
||||
location: str = ""
|
||||
name: str = ""
|
||||
p12_file: dict[str, Any] | None = None
|
||||
page_number: float = 1
|
||||
password: str = ""
|
||||
private_key_file: dict[str, Any] | None = None
|
||||
reason: str = ""
|
||||
show_logo: bool = True
|
||||
show_signature: bool = False
|
||||
sign_mode: Literal["MANUAL", "AUTO"] = "MANUAL"
|
||||
|
||||
|
||||
class ChangeMetadataParams(ApiModel):
|
||||
author: str = ""
|
||||
creation_date: str | None = None
|
||||
creator: str = ""
|
||||
custom_metadata: list[dict[str, str]] = []
|
||||
delete_all: bool = False
|
||||
keywords: str = ""
|
||||
modification_date: str | None = None
|
||||
producer: str = ""
|
||||
subject: str = ""
|
||||
title: str = ""
|
||||
trapped: Literal["True", "False", "Unknown"] | None = None
|
||||
|
||||
|
||||
class ChangePermissionsParams(ApiModel):
|
||||
prevent_assembly: bool = False
|
||||
prevent_extract_content: bool = False
|
||||
prevent_extract_for_accessibility: bool = False
|
||||
prevent_fill_in_form: bool = False
|
||||
prevent_modify: bool = False
|
||||
prevent_modify_annotations: bool = False
|
||||
prevent_printing: bool = False
|
||||
prevent_printing_faithful: bool = False
|
||||
|
||||
|
||||
class AddPasswordParams(ApiModel):
|
||||
key_length: float = 128
|
||||
owner_password: str = ""
|
||||
password: str = ""
|
||||
permissions: ChangePermissionsParams = ChangePermissionsParams.model_validate(
|
||||
{
|
||||
"preventAssembly": False,
|
||||
"preventExtractContent": False,
|
||||
"preventExtractForAccessibility": False,
|
||||
"preventFillInForm": False,
|
||||
"preventModify": False,
|
||||
"preventModifyAnnotations": False,
|
||||
"preventPrinting": False,
|
||||
"preventPrintingFaithful": False,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class CompressParams(ApiModel):
|
||||
compression_level: float = 5
|
||||
compression_method: Literal["quality", "filesize"] = "quality"
|
||||
expected_size: str = ""
|
||||
file_size_unit: Literal["KB", "MB"] = "MB"
|
||||
file_size_value: str = ""
|
||||
grayscale: bool = False
|
||||
|
||||
|
||||
class ConvertParams(ApiModel):
|
||||
cbz_options: dict[str, bool] = {"optimizeForEbook": False}
|
||||
cbz_output_options: dict[str, float] = {"dpi": 150}
|
||||
email_options: dict[str, Any] = {
|
||||
"includeAttachments": True,
|
||||
"maxAttachmentSizeMB": 10,
|
||||
"downloadHtml": False,
|
||||
"includeAllRecipients": False,
|
||||
}
|
||||
from_extension: str = ""
|
||||
html_options: dict[str, float] = {"zoomLevel": 1}
|
||||
image_options: dict[str, Any] = {
|
||||
"colorType": "color",
|
||||
"dpi": 300,
|
||||
"singleOrMultiple": "multiple",
|
||||
"fitOption": "maintainAspectRatio",
|
||||
"autoRotate": True,
|
||||
"combineImages": True,
|
||||
}
|
||||
is_smart_detection: bool = False
|
||||
pdfa_options: dict[str, str] = {"outputFormat": "pdfa-1"}
|
||||
smart_detection_type: Literal["mixed", "images", "web", "none"] = "none"
|
||||
to_extension: str = ""
|
||||
|
||||
|
||||
class CropParams(ApiModel):
|
||||
crop_area: dict[str, float] = {"x": 0, "y": 0, "width": 595, "height": 842}
|
||||
|
||||
|
||||
class EditTableOfContentsParams(ApiModel):
|
||||
bookmarks: list[dict[str, Any]] = []
|
||||
replace_existing: bool = True
|
||||
|
||||
|
||||
class ExtractImagesParams(ApiModel):
|
||||
allow_duplicates: bool = False
|
||||
format: Literal["png", "jpg", "gif"] = "png"
|
||||
|
||||
|
||||
class ExtractPagesParams(ApiModel):
|
||||
page_numbers: str = ""
|
||||
|
||||
|
||||
class FlattenParams(ApiModel):
|
||||
flatten_only_forms: bool = False
|
||||
|
||||
|
||||
class MergeParams(ApiModel):
|
||||
generate_table_of_contents: bool = False
|
||||
remove_digital_signature: bool = False
|
||||
|
||||
|
||||
class OcrParams(ApiModel):
|
||||
additional_options: list[str] = []
|
||||
languages: list[str] = []
|
||||
ocr_render_type: str = "hocr"
|
||||
ocr_type: str = "skip-text"
|
||||
|
||||
|
||||
class OverlayPdfsParams(ApiModel):
|
||||
counts: list[float] = []
|
||||
overlay_files: list[dict[str, Any]] = []
|
||||
overlay_mode: Literal["SequentialOverlay", "InterleavedOverlay", "FixedRepeatOverlay"] = "SequentialOverlay"
|
||||
overlay_position: Literal[0, 1] = 0
|
||||
|
||||
|
||||
class PageLayoutParams(ApiModel):
|
||||
add_border: bool = False
|
||||
pages_per_sheet: float = 4
|
||||
|
||||
|
||||
class PdfToSinglePageParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class RedactParams(ApiModel):
|
||||
convert_pdfto_image: bool = True
|
||||
custom_padding: float = 0.1
|
||||
mode: Literal["automatic", "manual"] = "automatic"
|
||||
redact_color: str = "#000000"
|
||||
use_regex: bool = False
|
||||
whole_word_search: bool = False
|
||||
words_to_redact: list[str] = []
|
||||
|
||||
|
||||
class RemoveAnnotationsParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class RemoveBlanksParams(ApiModel):
|
||||
include_blank_pages: bool = False
|
||||
threshold: float = 10
|
||||
white_percent: float = 99.9
|
||||
|
||||
|
||||
class RemoveCertSignParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class RemoveImageParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class RemovePagesParams(ApiModel):
|
||||
page_numbers: str = ""
|
||||
|
||||
|
||||
class RemovePasswordParams(ApiModel):
|
||||
password: str = ""
|
||||
|
||||
|
||||
class ReorganizePagesParams(ApiModel):
|
||||
custom_mode: str | None = None
|
||||
page_numbers: str | None = None
|
||||
|
||||
|
||||
class RepairParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class ReplaceColorParams(ApiModel):
|
||||
back_ground_color: str = "#ffffff"
|
||||
high_contrast_color_combination: Literal[
|
||||
"WHITE_TEXT_ON_BLACK", "BLACK_TEXT_ON_WHITE", "YELLOW_TEXT_ON_BLACK", "GREEN_TEXT_ON_BLACK"
|
||||
] = "WHITE_TEXT_ON_BLACK"
|
||||
replace_and_invert_option: Literal[
|
||||
"HIGH_CONTRAST_COLOR", "CUSTOM_COLOR", "FULL_INVERSION", "COLOR_SPACE_CONVERSION"
|
||||
] = "HIGH_CONTRAST_COLOR"
|
||||
text_color: str = "#000000"
|
||||
|
||||
|
||||
class RotateParams(ApiModel):
|
||||
angle: float = 0
|
||||
|
||||
|
||||
class SanitizeParams(ApiModel):
|
||||
remove_embedded_files: bool = True
|
||||
remove_fonts: bool = False
|
||||
remove_java_script: bool = True
|
||||
remove_links: bool = False
|
||||
remove_metadata: bool = False
|
||||
remove_xmpmetadata: bool = False
|
||||
|
||||
|
||||
class ScalePagesParams(ApiModel):
|
||||
page_size: Literal["KEEP", "A0", "A1", "A2", "A3", "A4", "A5", "A6", "LETTER", "LEGAL"] | None = None
|
||||
scale_factor: float = 1
|
||||
|
||||
|
||||
class ScannerImageSplitParams(ApiModel):
|
||||
angle_threshold: float = 10
|
||||
border_size: float = 1
|
||||
min_area: float = 10000
|
||||
min_contour_area: float = 500
|
||||
tolerance: float = 30
|
||||
|
||||
|
||||
class SignParams(ApiModel):
|
||||
font_family: str = "Helvetica"
|
||||
font_size: float = 16
|
||||
location: str = "Digital"
|
||||
reason: str = "Document signing"
|
||||
signature_data: str | None = None
|
||||
signature_position: dict[str, float] | None = None
|
||||
signature_type: Literal["text", "image", "canvas"] = "canvas"
|
||||
signer_name: str = ""
|
||||
text_color: str = "#000000"
|
||||
|
||||
|
||||
class SplitParams(ApiModel):
|
||||
allow_duplicates: bool = False
|
||||
bookmark_level: str = "1"
|
||||
duplex_mode: bool = False
|
||||
h_div: str = "2"
|
||||
include_metadata: bool = False
|
||||
merge: bool = False
|
||||
method: Literal[
|
||||
"", "byPages", "bySections", "bySize", "byPageCount", "byDocCount", "byChapters", "byPageDivider"
|
||||
] = ""
|
||||
pages: str = ""
|
||||
split_value: str = ""
|
||||
v_div: str = "2"
|
||||
|
||||
|
||||
class UnlockPdfformsParams(ApiModel):
|
||||
pass
|
||||
|
||||
|
||||
class WatermarkParams(ApiModel):
|
||||
alphabet: str = "roman"
|
||||
convert_pdfto_image: bool = False
|
||||
custom_color: str = "#d3d3d3"
|
||||
font_size: float = 12
|
||||
height_spacer: float = 50
|
||||
opacity: float = 50
|
||||
rotation: float = 0
|
||||
watermark_image: dict[str, Any] | None = None
|
||||
watermark_text: str = ""
|
||||
watermark_type: Literal["text", "image"] | None = None
|
||||
width_spacer: float = 50
|
||||
|
||||
|
||||
type ParamToolModel = (
|
||||
AddAttachmentsParams
|
||||
| AddPasswordParams
|
||||
| AdjustContrastParams
|
||||
| AutoRenameParams
|
||||
| AutomateParams
|
||||
| BookletImpositionParams
|
||||
| CertSignParams
|
||||
| ChangeMetadataParams
|
||||
| ChangePermissionsParams
|
||||
| CompressParams
|
||||
| ConvertParams
|
||||
| CropParams
|
||||
| EditTableOfContentsParams
|
||||
| ExtractImagesParams
|
||||
| ExtractPagesParams
|
||||
| FlattenParams
|
||||
| MergeParams
|
||||
| OcrParams
|
||||
| OverlayPdfsParams
|
||||
| PageLayoutParams
|
||||
| PdfToSinglePageParams
|
||||
| RedactParams
|
||||
| RemoveAnnotationsParams
|
||||
| RemoveBlanksParams
|
||||
| RemoveCertSignParams
|
||||
| RemoveImageParams
|
||||
| RemovePagesParams
|
||||
| RemovePasswordParams
|
||||
| ReorganizePagesParams
|
||||
| RepairParams
|
||||
| ReplaceColorParams
|
||||
| RotateParams
|
||||
| SanitizeParams
|
||||
| ScalePagesParams
|
||||
| ScannerImageSplitParams
|
||||
| SignParams
|
||||
| SplitParams
|
||||
| UnlockPdfformsParams
|
||||
| WatermarkParams
|
||||
)
|
||||
type ParamToolModelType = type[ParamToolModel]
|
||||
|
||||
|
||||
class OperationId(StrEnum):
|
||||
ADD_ATTACHMENTS = "addAttachments"
|
||||
ADD_PASSWORD = "addPassword"
|
||||
ADJUST_CONTRAST = "adjustContrast"
|
||||
AUTO_RENAME = "autoRename"
|
||||
AUTOMATE = "automate"
|
||||
BOOKLET_IMPOSITION = "bookletImposition"
|
||||
CERT_SIGN = "certSign"
|
||||
CHANGE_METADATA = "changeMetadata"
|
||||
CHANGE_PERMISSIONS = "changePermissions"
|
||||
COMPRESS = "compress"
|
||||
CONVERT = "convert"
|
||||
CROP = "crop"
|
||||
EDIT_TABLE_OF_CONTENTS = "editTableOfContents"
|
||||
EXTRACT_IMAGES = "extractImages"
|
||||
EXTRACT_PAGES = "extractPages"
|
||||
FLATTEN = "flatten"
|
||||
MERGE = "merge"
|
||||
OCR = "ocr"
|
||||
OVERLAY_PDFS = "overlayPdfs"
|
||||
PAGE_LAYOUT = "pageLayout"
|
||||
PDF_TO_SINGLE_PAGE = "pdfToSinglePage"
|
||||
REDACT = "redact"
|
||||
REMOVE_ANNOTATIONS = "removeAnnotations"
|
||||
REMOVE_BLANKS = "removeBlanks"
|
||||
REMOVE_CERT_SIGN = "removeCertSign"
|
||||
REMOVE_IMAGE = "removeImage"
|
||||
REMOVE_PAGES = "removePages"
|
||||
REMOVE_PASSWORD = "removePassword"
|
||||
REORGANIZE_PAGES = "reorganizePages"
|
||||
REPAIR = "repair"
|
||||
REPLACE_COLOR = "replaceColor"
|
||||
ROTATE = "rotate"
|
||||
SANITIZE = "sanitize"
|
||||
SCALE_PAGES = "scalePages"
|
||||
SCANNER_IMAGE_SPLIT = "scannerImageSplit"
|
||||
SIGN = "sign"
|
||||
SPLIT = "split"
|
||||
UNLOCK_PDFFORMS = "unlockPDFForms"
|
||||
WATERMARK = "watermark"
|
||||
|
||||
|
||||
OPERATIONS: dict[OperationId, ParamToolModelType] = {
|
||||
OperationId.ADD_ATTACHMENTS: AddAttachmentsParams,
|
||||
OperationId.ADD_PASSWORD: AddPasswordParams,
|
||||
OperationId.ADJUST_CONTRAST: AdjustContrastParams,
|
||||
OperationId.AUTO_RENAME: AutoRenameParams,
|
||||
OperationId.AUTOMATE: AutomateParams,
|
||||
OperationId.BOOKLET_IMPOSITION: BookletImpositionParams,
|
||||
OperationId.CERT_SIGN: CertSignParams,
|
||||
OperationId.CHANGE_METADATA: ChangeMetadataParams,
|
||||
OperationId.CHANGE_PERMISSIONS: ChangePermissionsParams,
|
||||
OperationId.COMPRESS: CompressParams,
|
||||
OperationId.CONVERT: ConvertParams,
|
||||
OperationId.CROP: CropParams,
|
||||
OperationId.EDIT_TABLE_OF_CONTENTS: EditTableOfContentsParams,
|
||||
OperationId.EXTRACT_IMAGES: ExtractImagesParams,
|
||||
OperationId.EXTRACT_PAGES: ExtractPagesParams,
|
||||
OperationId.FLATTEN: FlattenParams,
|
||||
OperationId.MERGE: MergeParams,
|
||||
OperationId.OCR: OcrParams,
|
||||
OperationId.OVERLAY_PDFS: OverlayPdfsParams,
|
||||
OperationId.PAGE_LAYOUT: PageLayoutParams,
|
||||
OperationId.PDF_TO_SINGLE_PAGE: PdfToSinglePageParams,
|
||||
OperationId.REDACT: RedactParams,
|
||||
OperationId.REMOVE_ANNOTATIONS: RemoveAnnotationsParams,
|
||||
OperationId.REMOVE_BLANKS: RemoveBlanksParams,
|
||||
OperationId.REMOVE_CERT_SIGN: RemoveCertSignParams,
|
||||
OperationId.REMOVE_IMAGE: RemoveImageParams,
|
||||
OperationId.REMOVE_PAGES: RemovePagesParams,
|
||||
OperationId.REMOVE_PASSWORD: RemovePasswordParams,
|
||||
OperationId.REORGANIZE_PAGES: ReorganizePagesParams,
|
||||
OperationId.REPAIR: RepairParams,
|
||||
OperationId.REPLACE_COLOR: ReplaceColorParams,
|
||||
OperationId.ROTATE: RotateParams,
|
||||
OperationId.SANITIZE: SanitizeParams,
|
||||
OperationId.SCALE_PAGES: ScalePagesParams,
|
||||
OperationId.SCANNER_IMAGE_SPLIT: ScannerImageSplitParams,
|
||||
OperationId.SIGN: SignParams,
|
||||
OperationId.SPLIT: SplitParams,
|
||||
OperationId.UNLOCK_PDFFORMS: UnlockPdfformsParams,
|
||||
OperationId.WATERMARK: WatermarkParams,
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Shared services used by the Stirling AI runtime."""
|
||||
|
||||
from .runtime import AppRuntime, build_model_settings, build_runtime
|
||||
|
||||
__all__ = [
|
||||
"AppRuntime",
|
||||
"build_model_settings",
|
||||
"build_runtime",
|
||||
]
|
||||
@@ -0,0 +1,51 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai.models import Model, infer_model
|
||||
from pydantic_ai.settings import ModelSettings
|
||||
|
||||
from stirling.config import AppSettings
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AppRuntime:
|
||||
settings: AppSettings
|
||||
fast_model: Model
|
||||
smart_model: Model
|
||||
|
||||
@property
|
||||
def fast_model_settings(self) -> ModelSettings:
|
||||
return build_model_settings(self.settings.fast_model_max_tokens)
|
||||
|
||||
@property
|
||||
def smart_model_settings(self) -> ModelSettings:
|
||||
return build_model_settings(self.settings.smart_model_max_tokens)
|
||||
|
||||
|
||||
def build_model_settings(max_tokens: int | None) -> ModelSettings:
|
||||
model_settings: ModelSettings = {}
|
||||
if max_tokens is not None:
|
||||
model_settings["max_tokens"] = max_tokens
|
||||
return model_settings
|
||||
|
||||
|
||||
def validate_structured_output_support(model: Model, model_name: str) -> None:
|
||||
# Pydantic AI's dedicated test model does not advertise native structured output,
|
||||
# but we still use it in unit tests as a non-production stand-in.
|
||||
if model_name == "test":
|
||||
return
|
||||
if not model.profile.supports_json_schema_output:
|
||||
raise ValueError(f"Unsupported model {model_name}. This model does not support structured outputs.")
|
||||
|
||||
|
||||
def build_runtime(settings: AppSettings) -> AppRuntime:
|
||||
fast_model = infer_model(settings.fast_model_name)
|
||||
smart_model = infer_model(settings.smart_model_name)
|
||||
validate_structured_output_support(fast_model, settings.fast_model_name)
|
||||
validate_structured_output_support(smart_model, settings.smart_model_name)
|
||||
return AppRuntime(
|
||||
settings=settings,
|
||||
fast_model=fast_model,
|
||||
smart_model=smart_model,
|
||||
)
|
||||
Reference in New Issue
Block a user