mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 06:12:49 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
283 lines
12 KiB
Python
283 lines
12 KiB
Python
"""Agent management sharing functionality."""
|
|
|
|
import datetime
|
|
import secrets
|
|
|
|
from flask import current_app, jsonify, make_response, request
|
|
from flask_restx import fields, Namespace, Resource
|
|
from sqlalchemy import text as _sql_text
|
|
|
|
from docsgpt.api import api
|
|
from docsgpt.core.settings import settings
|
|
from docsgpt.api.user.base import resolve_tool_details
|
|
from docsgpt.storage.db.base_repository import looks_like_uuid
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
from docsgpt.storage.db.repositories.users import UsersRepository
|
|
from docsgpt.storage.db.session import db_readonly, db_session
|
|
from docsgpt.utils import generate_image_url
|
|
|
|
agents_sharing_ns = Namespace(
|
|
"agents", description="Agent management operations", path="/api"
|
|
)
|
|
|
|
|
|
def _serialize_agent_basic(agent: dict) -> dict:
|
|
"""Shape a PG agent row into the API response dict."""
|
|
source_id = agent.get("source_id")
|
|
return {
|
|
"id": str(agent["id"]),
|
|
"user": agent.get("user_id", ""),
|
|
"name": agent.get("name", ""),
|
|
"image": (
|
|
generate_image_url(
|
|
agent["image"], agent["id"], agent.get("user_id")
|
|
)
|
|
if agent.get("image")
|
|
else ""
|
|
),
|
|
"description": agent.get("description", ""),
|
|
"source": str(source_id) if source_id else "",
|
|
"chunks": str(agent["chunks"]) if agent.get("chunks") is not None else "0",
|
|
"retriever": agent.get("retriever", "classic") or "classic",
|
|
"prompt_id": str(agent["prompt_id"]) if agent.get("prompt_id") else "default",
|
|
"tools": agent.get("tools", []) or [],
|
|
"tool_details": resolve_tool_details(agent.get("tools", []) or []),
|
|
"agent_type": agent.get("agent_type", "") or "",
|
|
"status": agent.get("status", "") or "",
|
|
"json_schema": agent.get("json_schema"),
|
|
"limited_token_mode": agent.get("limited_token_mode", False),
|
|
"token_limit": agent.get("token_limit") or settings.DEFAULT_AGENT_LIMITS["token_limit"],
|
|
"limited_request_mode": agent.get("limited_request_mode", False),
|
|
"request_limit": agent.get("request_limit") or settings.DEFAULT_AGENT_LIMITS["request_limit"],
|
|
"created_at": agent.get("created_at", ""),
|
|
"updated_at": agent.get("updated_at", ""),
|
|
"shared": bool(agent.get("shared", False)),
|
|
"shared_token": agent.get("shared_token", "") or "",
|
|
"shared_metadata": agent.get("shared_metadata", {}) or {},
|
|
}
|
|
|
|
|
|
@agents_sharing_ns.route("/shared_agent")
|
|
class SharedAgent(Resource):
|
|
@api.doc(
|
|
params={
|
|
"token": "Shared token of the agent",
|
|
},
|
|
description="Get a shared agent by token or ID",
|
|
)
|
|
def get(self):
|
|
shared_token = request.args.get("token")
|
|
|
|
if not shared_token:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Token or ID is required"}), 400
|
|
)
|
|
try:
|
|
with db_readonly() as conn:
|
|
shared_agent = AgentsRepository(conn).find_by_shared_token(
|
|
shared_token,
|
|
)
|
|
if not shared_agent:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Shared agent not found"}),
|
|
404,
|
|
)
|
|
agent_id = str(shared_agent["id"])
|
|
data = _serialize_agent_basic(shared_agent)
|
|
|
|
if data["tools"]:
|
|
enriched_tools = []
|
|
for detail in data["tool_details"]:
|
|
enriched_tools.append(detail.get("name", ""))
|
|
data["tools"] = enriched_tools
|
|
decoded_token = getattr(request, "decoded_token", None)
|
|
if decoded_token:
|
|
user_id = decoded_token.get("sub")
|
|
owner_id = shared_agent.get("user_id")
|
|
|
|
if user_id != owner_id:
|
|
with db_session() as conn:
|
|
users_repo = UsersRepository(conn)
|
|
users_repo.upsert(user_id)
|
|
users_repo.add_shared(user_id, agent_id)
|
|
return make_response(jsonify(data), 200)
|
|
except Exception as err:
|
|
current_app.logger.error(f"Error retrieving shared agent: {err}")
|
|
return make_response(jsonify({"success": False}), 400)
|
|
|
|
|
|
@agents_sharing_ns.route("/shared_agents")
|
|
class SharedAgents(Resource):
|
|
@api.doc(description="Get shared agents explicitly shared with the user")
|
|
def get(self):
|
|
try:
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user_id = decoded_token.get("sub")
|
|
|
|
with db_session() as conn:
|
|
users_repo = UsersRepository(conn)
|
|
user_doc = users_repo.upsert(user_id)
|
|
shared_with_ids = (
|
|
user_doc.get("agent_preferences", {}).get("shared_with_me", [])
|
|
if isinstance(user_doc.get("agent_preferences"), dict)
|
|
else []
|
|
)
|
|
# Keep only UUID-shaped ids; ObjectId leftovers are stripped below.
|
|
uuid_ids = [sid for sid in shared_with_ids if looks_like_uuid(sid)]
|
|
non_uuid_ids = [sid for sid in shared_with_ids if not looks_like_uuid(sid)]
|
|
|
|
if uuid_ids:
|
|
result = conn.execute(
|
|
_sql_text(
|
|
"SELECT * FROM agents "
|
|
"WHERE id = ANY(CAST(:ids AS uuid[])) "
|
|
"AND shared = true"
|
|
),
|
|
{"ids": uuid_ids},
|
|
)
|
|
shared_agents = [dict(row._mapping) for row in result.fetchall()]
|
|
else:
|
|
shared_agents = []
|
|
|
|
found_ids_set = {str(agent["id"]) for agent in shared_agents}
|
|
stale_ids = [sid for sid in uuid_ids if sid not in found_ids_set]
|
|
stale_ids.extend(non_uuid_ids)
|
|
if stale_ids:
|
|
users_repo.remove_shared_bulk(user_id, stale_ids)
|
|
|
|
pinned_ids = set(
|
|
user_doc.get("agent_preferences", {}).get("pinned", [])
|
|
if isinstance(user_doc.get("agent_preferences"), dict)
|
|
else []
|
|
)
|
|
|
|
list_shared_agents = []
|
|
for agent in shared_agents:
|
|
agent_id_str = str(agent["id"])
|
|
list_shared_agents.append(
|
|
{
|
|
"id": agent_id_str,
|
|
"name": agent.get("name", ""),
|
|
"description": agent.get("description", ""),
|
|
"image": (
|
|
generate_image_url(
|
|
agent["image"], agent["id"], agent.get("user_id")
|
|
)
|
|
if agent.get("image")
|
|
else ""
|
|
),
|
|
"tools": agent.get("tools", []) or [],
|
|
"tool_details": resolve_tool_details(
|
|
agent.get("tools", []) or []
|
|
),
|
|
"agent_type": agent.get("agent_type", "") or "",
|
|
"status": agent.get("status", "") or "",
|
|
"json_schema": agent.get("json_schema"),
|
|
"limited_token_mode": agent.get("limited_token_mode", False),
|
|
"token_limit": agent.get("token_limit") or settings.DEFAULT_AGENT_LIMITS["token_limit"],
|
|
"limited_request_mode": agent.get("limited_request_mode", False),
|
|
"request_limit": agent.get("request_limit") or settings.DEFAULT_AGENT_LIMITS["request_limit"],
|
|
"created_at": agent.get("created_at", ""),
|
|
"updated_at": agent.get("updated_at", ""),
|
|
"pinned": agent_id_str in pinned_ids,
|
|
"shared": bool(agent.get("shared", False)),
|
|
"shared_token": agent.get("shared_token", "") or "",
|
|
"shared_metadata": agent.get("shared_metadata", {}) or {},
|
|
}
|
|
)
|
|
|
|
return make_response(jsonify(list_shared_agents), 200)
|
|
except Exception as err:
|
|
current_app.logger.error(f"Error retrieving shared agents: {err}")
|
|
return make_response(jsonify({"success": False}), 400)
|
|
|
|
|
|
@agents_sharing_ns.route("/share_agent")
|
|
class ShareAgent(Resource):
|
|
@api.expect(
|
|
api.model(
|
|
"ShareAgentModel",
|
|
{
|
|
"id": fields.String(required=True, description="ID of the agent"),
|
|
"shared": fields.Boolean(
|
|
required=True, description="Share or unshare the agent"
|
|
),
|
|
"username": fields.String(
|
|
required=False, description="Name of the user"
|
|
),
|
|
},
|
|
)
|
|
)
|
|
@api.doc(description="Share or unshare an agent")
|
|
def put(self):
|
|
decoded_token = request.decoded_token
|
|
if not decoded_token:
|
|
return make_response(jsonify({"success": False}), 401)
|
|
user = decoded_token.get("sub")
|
|
|
|
data = request.get_json()
|
|
if not data:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Missing JSON body"}), 400
|
|
)
|
|
agent_id = data.get("id")
|
|
shared = data.get("shared")
|
|
username = data.get("username", "")
|
|
|
|
if not agent_id:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "ID is required"}), 400
|
|
)
|
|
if shared is None:
|
|
return make_response(
|
|
jsonify(
|
|
{
|
|
"success": False,
|
|
"message": "Shared parameter is required and must be true or false",
|
|
}
|
|
),
|
|
400,
|
|
)
|
|
shared_token = None
|
|
try:
|
|
with db_session() as conn:
|
|
repo = AgentsRepository(conn)
|
|
agent = repo.get_any(agent_id, user)
|
|
if not agent:
|
|
return make_response(
|
|
jsonify({"success": False, "message": "Agent not found"}), 404
|
|
)
|
|
if shared:
|
|
shared_metadata = {
|
|
"shared_by": username,
|
|
"shared_at": datetime.datetime.now(
|
|
datetime.timezone.utc
|
|
).isoformat(),
|
|
}
|
|
shared_token = secrets.token_urlsafe(32)
|
|
repo.update(
|
|
str(agent["id"]), user,
|
|
{
|
|
"shared": True,
|
|
"shared_token": shared_token,
|
|
"shared_metadata": shared_metadata,
|
|
},
|
|
)
|
|
else:
|
|
repo.update(
|
|
str(agent["id"]), user,
|
|
{
|
|
"shared": False,
|
|
"shared_token": None,
|
|
"shared_metadata": None,
|
|
},
|
|
)
|
|
except Exception as err:
|
|
current_app.logger.error(f"Error sharing/unsharing agent: {err}", exc_info=True)
|
|
return make_response(jsonify({"success": False, "error": "Failed to update agent sharing status"}), 400)
|
|
return make_response(
|
|
jsonify({"success": True, "shared_token": shared_token}), 200
|
|
)
|