TimHoogervorst/Agents
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added quick connect auth from jellyfin, still needs to have some more cleaning before push to prod

Browse Source
This commit is contained in:
TimHoogervorst
2026-05-25 10:19:50 +02:00
parent 3cd2e4dfbb
commit bf358f7248
16 changed files with 1318 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/auth_store.py
+316
View File
@@ -0,0 +1,316 @@
"""
Auth Store — SQLite-backed persistence for Discord-to-service authentication.
Two tables:
- link_tokens : one-time tokens sent via Discord DM to initiate login
- user_auth : per-user, per-service credentials (Jellyfin token, etc.)
Thread-safe via WAL mode and a shared lock. No passwords are ever stored
— only opaque service tokens (e.g. Jellyfin AccessToken).
"""
from __future__ import annotations
import logging
import secrets
import sqlite3
import threading
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Optional
from core.config import get_config
logger = logging.getLogger("auth_store")
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
AUTH_DB_PATH = get_config("AUTH_DB_PATH", "data/auth.db")
TOKEN_EXPIRY_MINUTES = int(get_config("AUTH_TOKEN_EXPIRY", "10"))
# ---------------------------------------------------------------------------
# Singleton handle
# ---------------------------------------------------------------------------
_db_path: Path | None = None
_db_lock = threading.Lock()
def _resolve_path() -> Path:
"""Turn AUTH_DB_PATH into an absolute path, creating parent dirs."""
global _db_path
if _db_path is not None:
return _db_path
p = Path(AUTH_DB_PATH)
if not p.is_absolute():
# Relative to the project root (two levels above this file)
project_root = Path(__file__).resolve().parent.parent
p = project_root / p
p.parent.mkdir(parents=True, exist_ok=True)
_db_path = p
return p
def _get_conn() -> sqlite3.Connection:
"""Return a thread-local connection to the auth database."""
import sqlite3
conn = sqlite3.connect(str(_resolve_path()), check_same_thread=False)
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA foreign_keys=ON")
conn.row_factory = sqlite3.Row
return conn
# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------
_SCHEMA = """
CREATE TABLE IF NOT EXISTS link_tokens (
token TEXT PRIMARY KEY,
discord_user_id INTEGER NOT NULL,
service TEXT NOT NULL,
expires_at TEXT NOT NULL,
used INTEGER DEFAULT 0,
created_at TEXT DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS user_auth (
discord_user_id INTEGER NOT NULL,
service TEXT NOT NULL,
external_user_id TEXT,
external_name TEXT,
credentials TEXT,
linked_at TEXT DEFAULT (datetime('now')),
is_active INTEGER DEFAULT 1,
PRIMARY KEY (discord_user_id, service)
);
"""
_initialized = False
def _ensure_schema() -> None:
global _initialized
if _initialized:
return
with _db_lock:
if _initialized:
return
conn = _get_conn()
conn.executescript(_SCHEMA)
conn.commit()
conn.close()
_initialized = True
logger.info("Auth store initialized at %s", _resolve_path())
# ---------------------------------------------------------------------------
# Public API — Link Tokens
# ---------------------------------------------------------------------------
def create_token(discord_user_id: int, service: str) -> str:
"""Generate a one-time link token. Expires after TOKEN_EXPIRY_MINUTES."""
_ensure_schema()
token = secrets.token_urlsafe(32)
expires = (datetime.now(timezone.utc) + timedelta(minutes=TOKEN_EXPIRY_MINUTES)).isoformat()
with _db_lock:
conn = _get_conn()
conn.execute(
"INSERT INTO link_tokens (token, discord_user_id, service, expires_at) VALUES (?, ?, ?, ?)",
(token, discord_user_id, service, expires),
)
conn.commit()
conn.close()
logger.info("Created link token for user %s / service %s", discord_user_id, service)
return token
def validate_token(token: str) -> tuple[int, str] | None:
"""Read-only validation — does NOT consume the token.
Returns (discord_user_id, service) if the token exists, is unused,
and has not expired. Returns None otherwise.
"""
_ensure_schema()
with _db_lock:
conn = _get_conn()
row = conn.execute(
"SELECT discord_user_id, service, used, expires_at FROM link_tokens WHERE token = ?",
(token,),
).fetchone()
conn.close()
if row is None:
return None
if row["used"]:
return None
expires = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires:
return None
return (row["discord_user_id"], row["service"])
def consume_token(token: str) -> tuple[int, str] | None:
"""Validate and consume a link token. Returns (discord_user_id, service) or None.
A token is valid if:
- It exists
- It has not been used
- It has not expired
"""
_ensure_schema()
with _db_lock:
conn = _get_conn()
row = conn.execute(
"SELECT discord_user_id, service, used, expires_at FROM link_tokens WHERE token = ?",
(token,),
).fetchone()
if row is None:
conn.close()
return None
if row["used"]:
conn.close()
logger.warning("Token already used: %s…", token[:8])
return None
expires = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires:
conn.close()
logger.warning("Token expired: %s…", token[:8])
return None
conn.execute("UPDATE link_tokens SET used = 1 WHERE token = ?", (token,))
conn.commit()
result = (row["discord_user_id"], row["service"])
conn.close()
logger.info("Token consumed: %s… → user=%s service=%s", token[:8], result[0], result[1])
return result
# ---------------------------------------------------------------------------
# Public API — User Auth
# ---------------------------------------------------------------------------
def store_auth(
discord_user_id: int,
service: str,
*,
external_user_id: str = "",
external_name: str = "",
credentials: dict | None = None,
) -> None:
"""Store or update authentication for a user on a service."""
_ensure_schema()
import json
creds_json = json.dumps(credentials) if credentials else "{}"
with _db_lock:
conn = _get_conn()
conn.execute(
"""INSERT INTO user_auth (discord_user_id, service, external_user_id, external_name, credentials, linked_at)
VALUES (?, ?, ?, ?, ?, datetime('now'))
ON CONFLICT(discord_user_id, service) DO UPDATE SET
external_user_id = excluded.external_user_id,
external_name = excluded.external_name,
credentials = excluded.credentials,
linked_at = datetime('now'),
is_active = 1""",
(discord_user_id, service, external_user_id, external_name, creds_json),
)
conn.commit()
conn.close()
logger.info("Stored auth for user %s on %s as %s", discord_user_id, service, external_name)
def get_auth(discord_user_id: int, service: str) -> dict | None:
"""Retrieve stored auth for a user on a service. Returns None if not linked."""
_ensure_schema()
import json
with _db_lock:
conn = _get_conn()
row = conn.execute(
"SELECT * FROM user_auth WHERE discord_user_id = ? AND service = ? AND is_active = 1",
(discord_user_id, service),
).fetchone()
conn.close()
if row is None:
return None
credentials = json.loads(row["credentials"]) if row["credentials"] else {}
return {
"discord_user_id": row["discord_user_id"],
"service": row["service"],
"external_user_id": row["external_user_id"],
"external_name": row["external_name"],
"credentials": credentials,
"linked_at": row["linked_at"],
}
def is_authenticated(discord_user_id: int, service: str) -> bool:
"""Quick check: is this user linked to this service?"""
return get_auth(discord_user_id, service) is not None
def list_services(discord_user_id: int) -> list[str]:
"""Return list of service names this user has linked."""
_ensure_schema()
with _db_lock:
conn = _get_conn()
rows = conn.execute(
"SELECT service FROM user_auth WHERE discord_user_id = ? AND is_active = 1",
(discord_user_id,),
).fetchall()
conn.close()
return [r["service"] for r in rows]
def revoke(discord_user_id: int, service: str) -> None:
"""Unlink a user from a service."""
_ensure_schema()
with _db_lock:
conn = _get_conn()
conn.execute(
"UPDATE user_auth SET is_active = 0 WHERE discord_user_id = ? AND service = ?",
(discord_user_id, service),
)
conn.commit()
conn.close()
logger.info("Revoked auth for user %s on %s", discord_user_id, service)
# ---------------------------------------------------------------------------
# Dev / testing — reset the entire store
# ---------------------------------------------------------------------------
def reset_all() -> None:
"""Truncate all auth tables — for development and testing only."""
_ensure_schema()
with _db_lock:
conn = _get_conn()
conn.execute("DELETE FROM link_tokens")
conn.execute("DELETE FROM user_auth")
conn.commit()
conn.close()
logger.warning("Auth store RESET — all tokens and auth records cleared.")
core/graph.py
+13 -29
View File
@@ -1,12 +1,13 @@
"""
LangGraph agent graph factory.
Builds a StateGraph that replaces the manual tool-calling loop in api/v1/chat.py.
The graph has two nodes:
Builds a StateGraph with two nodes:
- agent_node : calls the LLM (with system prompt + tool definitions)
- tool_node : executes tool calls via the existing skill system
A conditional edge routes tool_calls back to the agent, or ends the run.
When a tool fails due to missing authentication, the failure message is
relayed to the LLM, which tells the user to use /login.
"""
from __future__ import annotations
@@ -97,18 +98,14 @@ def _make_agent_node(
full: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
for m in messages:
if isinstance(m, dict):
# Already a plain dict — pass through.
# But fix tool_calls if they're in LangChain format.
d = dict(m)
tc = d.get("tool_calls")
if tc and isinstance(tc, list) and tc and isinstance(tc[0], dict) and "function" not in tc[0]:
d["tool_calls"] = _langchain_tc_to_openai(tc)
full.append(d)
else:
# LangChain message object → OpenAI-compatible dict
role = _lc_role_to_openai(getattr(m, "type", "user"))
d: dict[str, Any] = {"role": role, "content": getattr(m, "content", "")}
# Serialize tool_calls back to OpenAI format (if this is an AI msg)
tc = getattr(m, "tool_calls", None)
if tc:
d["tool_calls"] = _langchain_tc_to_openai(tc)
@@ -125,7 +122,6 @@ def _make_agent_node(
)
choice = resp.choices[0]
# Convert OpenAI tool_calls to the dict format LangChain expects.
raw_tool_calls = list(choice.message.tool_calls) if choice.message.tool_calls else []
tool_calls: list[dict[str, Any]] = []
for tc in raw_tool_calls:
@@ -153,9 +149,9 @@ def _make_tool_node(skill_names: list[str]):
"""
Return a callable that executes tool_calls from the last AI message.
This replaces LangGraph's built-in ToolNode — we call our own
`execute_tool()` pipeline so that skill-level auth, httpx sessions,
and ToolResult handling are fully preserved.
If a tool fails because the user isn't authenticated, the failure
message (which tells the user to /login) is returned to the LLM.
The LLM naturally relays the instructions to the user.
"""
async def tool_node(state: AgentState) -> dict[str, list]:
@@ -164,18 +160,16 @@ def _make_tool_node(skill_names: list[str]):
if not tool_calls:
return {"messages": []}
discord_user_id = state.get("discord_user_id")
results: list[ToolMessage] = []
for tc in tool_calls:
# Handle both LangChain format (top-level name/args) and
# OpenAI format (nested "function" key).
if isinstance(tc, dict):
if "function" in tc:
# OpenAI format: {"id":..., "function": {"name":..., "arguments":"..."}}
fn = tc["function"]
fn_name = fn.get("name", "")
fn_args_raw = fn.get("arguments", "{}")
else:
# LangChain format: {"name":..., "args":{...}, "id":...}
fn_name = tc.get("name", "")
fn_args_raw = tc.get("args", {})
tc_id = tc.get("id", "")
@@ -184,13 +178,15 @@ def _make_tool_node(skill_names: list[str]):
fn_args_raw = getattr(tc, "args", {})
tc_id = getattr(tc, "id", "")
# Parse args if they arrive as a JSON string
if isinstance(fn_args_raw, str):
fn_args = json.loads(fn_args_raw)
else:
fn_args = fn_args_raw
tr = await execute_tool(skill_names, fn_name, fn_args)
tr = await execute_tool(
skill_names, fn_name, fn_args,
discord_user_id=discord_user_id,
)
content = tr.content if tr else f"Tool '{fn_name}' is not available."
results.append(ToolMessage(content=content, tool_call_id=tc_id))
@@ -224,27 +220,16 @@ def create_agent_graph(
) -> StateGraph:
"""
Build and compile a LangGraph StateGraph for a single agent.
Parameters
----------
client : The OpenAI-compatible client (already authenticated).
agent_skills : Skill names assigned to the agent (e.g. ["seerr", "triage"]).
system_prompt : The fully-built system prompt (base + skill fragments).
model_name : Model identifier sent to the LLM provider.
Returns
-------
A compiled LangGraph graph ready for `.ainvoke()` or `.astream()`.
"""
tool_defs = get_all_tools(agent_skills)
graph = StateGraph(AgentState)
# Nodes
graph.add_node(
"agent_node",
_make_agent_node(client, system_prompt, tool_defs, model_name),
)
if tool_defs:
graph.add_node("tool_node", _make_tool_node(agent_skills))
graph.add_conditional_edges("agent_node", _should_continue, {
@@ -253,7 +238,6 @@ def create_agent_graph(
})
graph.add_edge("tool_node", "agent_node")
else:
# No tools — agent responds once and finishes
graph.add_edge("agent_node", END)
graph.set_entry_point("agent_node")
core/state.py
+1
View File
@@ -18,3 +18,4 @@ class AgentState(TypedDict):
"""
messages: Annotated[list, add_messages]
discord_user_id: int | None # set by the Discord bot, None for REST API calls