Add external resource registration
This commit is contained in:
49
core/api.py
49
core/api.py
@ -7,6 +7,7 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import parse_qsl, quote, urlsplit, urlunsplit
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, FastAPI, File, Form, HTTPException, Request, UploadFile
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
from starlette.responses import Response
|
||||
@ -86,6 +87,22 @@ class FlushMemoryRequest(BaseModel):
|
||||
project_id: str = "default"
|
||||
|
||||
|
||||
class ExternalResourceRequest(BaseModel):
|
||||
user_id: str = Field(min_length=1)
|
||||
user_key: str = Field(min_length=1)
|
||||
app_id: str = "default"
|
||||
project_id: str = "default"
|
||||
filename: str = Field(min_length=1)
|
||||
mime_type: str | None = None
|
||||
content_type: str | None = None
|
||||
size_bytes: int | None = Field(default=None, ge=0)
|
||||
sha256: str | None = None
|
||||
source_uri: str = Field(min_length=1)
|
||||
ingest_uri: str | None = None
|
||||
title: str | None = None
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class MemoryOverrideRequest(BaseModel):
|
||||
user_id: str = Field(min_length=1)
|
||||
user_key: str = Field(min_length=1)
|
||||
@ -196,6 +213,13 @@ def _body_for_log(body: bytes, content_type: str | None) -> Any:
|
||||
return {"content_type": content_type, "size_bytes": len(body)}
|
||||
|
||||
|
||||
def _backend_http_error_detail(exc: httpx.HTTPStatusError) -> Any:
|
||||
try:
|
||||
return exc.response.json()
|
||||
except ValueError:
|
||||
return exc.response.text
|
||||
|
||||
|
||||
def create_app(
|
||||
*,
|
||||
config: GatewayConfig | None = None,
|
||||
@ -366,6 +390,26 @@ def create_app(
|
||||
return {"resources": []}
|
||||
return {"resources": [resource]}
|
||||
|
||||
@router.post("/resources/external")
|
||||
async def register_external_resource(
|
||||
request: ExternalResourceRequest,
|
||||
) -> dict[str, Any]:
|
||||
require_user(request.user_id, request.user_key)
|
||||
return await service.register_external_resource(
|
||||
user_id=request.user_id,
|
||||
app_id=request.app_id,
|
||||
project_id=request.project_id,
|
||||
filename=request.filename,
|
||||
mime_type=request.mime_type,
|
||||
content_type=request.content_type,
|
||||
size_bytes=request.size_bytes,
|
||||
sha256=request.sha256,
|
||||
source_uri=request.source_uri,
|
||||
ingest_uri=request.ingest_uri,
|
||||
title=request.title,
|
||||
description=request.description,
|
||||
)
|
||||
|
||||
@router.delete("/resources/{resource_id}")
|
||||
async def delete_resource(
|
||||
resource_id: str,
|
||||
@ -412,6 +456,11 @@ def create_app(
|
||||
project_id=request.project_id,
|
||||
messages=[message.model_dump() for message in request.messages],
|
||||
)
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise HTTPException(
|
||||
status_code=exc.response.status_code,
|
||||
detail=_backend_http_error_detail(exc),
|
||||
) from exc
|
||||
except UploadTooLarge as exc:
|
||||
raise HTTPException(status_code=413, detail=str(exc)) from exc
|
||||
except InvalidAttachment as exc:
|
||||
|
||||
137
core/service.py
137
core/service.py
@ -51,6 +51,17 @@ def infer_content_type(filename: str | None, mime_type: str | None) -> str:
|
||||
return "doc"
|
||||
|
||||
|
||||
def normalize_content_type(
|
||||
filename: str | None,
|
||||
mime_type: str | None,
|
||||
content_type: str | None,
|
||||
) -> str:
|
||||
requested = (content_type or "").strip().lower()
|
||||
if requested in {"image", "audio", "pdf", "html", "text", "doc"}:
|
||||
return requested
|
||||
return infer_content_type(filename, mime_type or requested)
|
||||
|
||||
|
||||
def _safe_filename(filename: str | None) -> str:
|
||||
name = Path(filename or "upload.bin").name
|
||||
return name or "upload.bin"
|
||||
@ -301,6 +312,123 @@ class MemoryGatewayService:
|
||||
item["base64"] = base64.b64encode(content).decode("ascii")
|
||||
return item
|
||||
|
||||
async def register_external_resource(
|
||||
self,
|
||||
*,
|
||||
user_id: str,
|
||||
app_id: str,
|
||||
project_id: str,
|
||||
filename: str,
|
||||
mime_type: str | None,
|
||||
content_type: str | None,
|
||||
size_bytes: int | None,
|
||||
sha256: str | None,
|
||||
source_uri: str,
|
||||
ingest_uri: str | None,
|
||||
title: str | None,
|
||||
description: str | None,
|
||||
) -> dict[str, Any]:
|
||||
resource_id = new_resource_id()
|
||||
session_id = resource_session_id(user_id, resource_id)
|
||||
original_filename = _safe_filename(filename)
|
||||
normalized_content_type = normalize_content_type(
|
||||
original_filename,
|
||||
mime_type,
|
||||
content_type,
|
||||
)
|
||||
existing = None
|
||||
if sha256:
|
||||
existing = self.repository.find_active_resource_by_sha256(
|
||||
user_id=user_id,
|
||||
app_id=app_id,
|
||||
project_id=project_id,
|
||||
sha256=sha256,
|
||||
)
|
||||
if existing is not None:
|
||||
self._register_resource_attachment(existing, source="external_resource")
|
||||
return self._resource_summary(existing)
|
||||
|
||||
resource = self.repository.create_resource(
|
||||
id=resource_id,
|
||||
user_id=user_id,
|
||||
app_id=app_id,
|
||||
project_id=project_id,
|
||||
session_id=session_id,
|
||||
original_filename=original_filename,
|
||||
mime_type=mime_type,
|
||||
content_type=normalized_content_type,
|
||||
uri=source_uri,
|
||||
uri_public=False,
|
||||
sha256=sha256,
|
||||
size_bytes=size_bytes,
|
||||
title=title,
|
||||
description=description,
|
||||
status="ingesting",
|
||||
error_message=None,
|
||||
)
|
||||
self._register_resource_attachment(resource, source="external_resource")
|
||||
|
||||
try:
|
||||
await self._retry_backend_call(
|
||||
lambda: self.backend_client.add_memory(
|
||||
self._build_external_add_payload(
|
||||
resource=resource,
|
||||
user_id=user_id,
|
||||
app_id=app_id,
|
||||
project_id=project_id,
|
||||
filename=original_filename,
|
||||
ingest_uri=ingest_uri or source_uri,
|
||||
)
|
||||
)
|
||||
)
|
||||
await self._retry_backend_call(
|
||||
lambda: self.backend_client.flush_memory(session_id, app_id, project_id)
|
||||
)
|
||||
except Exception as exc:
|
||||
failed = self.repository.update_resource_status(
|
||||
resource_id,
|
||||
"failed",
|
||||
str(exc),
|
||||
)
|
||||
return self._resource_summary(failed or resource)
|
||||
|
||||
extracted = self.repository.update_resource_status(resource_id, "extracted")
|
||||
return self._resource_summary(extracted or resource)
|
||||
|
||||
def _build_external_add_payload(
|
||||
self,
|
||||
*,
|
||||
resource: dict[str, Any],
|
||||
user_id: str,
|
||||
app_id: str,
|
||||
project_id: str,
|
||||
filename: str,
|
||||
ingest_uri: str,
|
||||
) -> dict[str, Any]:
|
||||
content_item = {
|
||||
"type": str(resource["content_type"]),
|
||||
"name": filename,
|
||||
"ext": Path(filename).suffix.lstrip(".") or None,
|
||||
"uri": ingest_uri,
|
||||
"extras": {
|
||||
"resource_id": resource["id"],
|
||||
"source": "external_resource",
|
||||
},
|
||||
}
|
||||
return {
|
||||
"session_id": resource["session_id"],
|
||||
"app_id": app_id,
|
||||
"project_id": project_id,
|
||||
"messages": [
|
||||
{
|
||||
"sender_id": user_id,
|
||||
"role": "user",
|
||||
"timestamp": current_timestamp_ms(),
|
||||
"content": [content_item],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
def _resource_file_path(self, resource: dict[str, Any]) -> Path:
|
||||
uri = str(resource["uri"])
|
||||
parsed = urlparse(uri)
|
||||
@ -489,7 +617,12 @@ class MemoryGatewayService:
|
||||
raise
|
||||
return {"session_id": session_id, "backend": backend}
|
||||
|
||||
def _register_resource_attachment(self, resource: dict[str, Any]) -> None:
|
||||
def _register_resource_attachment(
|
||||
self,
|
||||
resource: dict[str, Any],
|
||||
*,
|
||||
source: str = "resource_upload",
|
||||
) -> None:
|
||||
self.repository.create_attachment(
|
||||
user_id=resource["user_id"],
|
||||
app_id=resource["app_id"],
|
||||
@ -499,7 +632,7 @@ class MemoryGatewayService:
|
||||
content_type=resource["content_type"],
|
||||
name=resource["original_filename"] or resource["id"],
|
||||
internal_uri=resource["uri"],
|
||||
source="resource_upload",
|
||||
source=source,
|
||||
sha256=resource["sha256"],
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user