remdb 0.3.114__py3-none-any.whl → 0.3.127__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of remdb might be problematic. Click here for more details.
- rem/agentic/agents/sse_simulator.py +2 -0
- rem/agentic/context.py +23 -3
- rem/agentic/mcp/tool_wrapper.py +29 -3
- rem/agentic/otel/setup.py +1 -0
- rem/agentic/providers/pydantic_ai.py +26 -2
- rem/api/main.py +4 -1
- rem/api/mcp_router/server.py +9 -3
- rem/api/mcp_router/tools.py +324 -2
- rem/api/routers/admin.py +218 -1
- rem/api/routers/chat/completions.py +250 -4
- rem/api/routers/chat/models.py +81 -7
- rem/api/routers/chat/otel_utils.py +33 -0
- rem/api/routers/chat/sse_events.py +17 -1
- rem/api/routers/chat/streaming.py +35 -1
- rem/api/routers/feedback.py +134 -14
- rem/api/routers/query.py +6 -3
- rem/cli/commands/README.md +42 -0
- rem/cli/commands/cluster.py +617 -168
- rem/cli/commands/configure.py +1 -3
- rem/cli/commands/db.py +66 -22
- rem/cli/commands/experiments.py +242 -26
- rem/cli/commands/schema.py +6 -5
- rem/config.py +8 -1
- rem/services/phoenix/client.py +59 -18
- rem/services/postgres/diff_service.py +108 -3
- rem/services/postgres/schema_generator.py +205 -4
- rem/services/session/compression.py +7 -0
- rem/settings.py +150 -18
- rem/sql/migrations/001_install.sql +156 -0
- rem/sql/migrations/002_install_models.sql +1864 -1
- rem/sql/migrations/004_cache_system.sql +548 -0
- rem/utils/__init__.py +18 -0
- rem/utils/schema_loader.py +94 -3
- rem/utils/sql_paths.py +146 -0
- rem/workers/__init__.py +3 -1
- rem/workers/db_listener.py +579 -0
- rem/workers/unlogged_maintainer.py +463 -0
- {remdb-0.3.114.dist-info → remdb-0.3.127.dist-info}/METADATA +213 -177
- {remdb-0.3.114.dist-info → remdb-0.3.127.dist-info}/RECORD +41 -36
- {remdb-0.3.114.dist-info → remdb-0.3.127.dist-info}/WHEEL +0 -0
- {remdb-0.3.114.dist-info → remdb-0.3.127.dist-info}/entry_points.txt +0 -0
|
@@ -265,6 +265,8 @@ async def stream_simulator_events(
|
|
|
265
265
|
message_id=message_id,
|
|
266
266
|
in_reply_to=in_reply_to,
|
|
267
267
|
session_id=session_id,
|
|
268
|
+
# Session info
|
|
269
|
+
session_name="SSE Demo Session",
|
|
268
270
|
# Quality indicators
|
|
269
271
|
confidence=0.95,
|
|
270
272
|
sources=["rem/api/routers/chat/sse_events.py", "rem/agentic/agents/sse_simulator.py"],
|
rem/agentic/context.py
CHANGED
|
@@ -2,10 +2,18 @@
|
|
|
2
2
|
Agent execution context and configuration.
|
|
3
3
|
|
|
4
4
|
Design pattern for session context that can be constructed from:
|
|
5
|
-
- HTTP headers (X-User-Id, X-Session-Id, X-Model-Name)
|
|
5
|
+
- HTTP headers (X-User-Id, X-Session-Id, X-Model-Name, X-Is-Eval, etc.)
|
|
6
6
|
- Direct instantiation for testing/CLI
|
|
7
7
|
|
|
8
|
-
|
|
8
|
+
Headers Mapping:
|
|
9
|
+
X-User-Id → context.user_id
|
|
10
|
+
X-Tenant-Id → context.tenant_id (default: "default")
|
|
11
|
+
X-Session-Id → context.session_id
|
|
12
|
+
X-Agent-Schema → context.agent_schema_uri (default: "rem")
|
|
13
|
+
X-Model-Name → context.default_model
|
|
14
|
+
X-Is-Eval → context.is_eval (marks session as evaluation)
|
|
15
|
+
|
|
16
|
+
Key Design Pattern:
|
|
9
17
|
- AgentContext is passed to agent factory, not stored in agents
|
|
10
18
|
- Enables session tracking across API, CLI, and test execution
|
|
11
19
|
- Supports header-based configuration override (model, schema URI)
|
|
@@ -66,6 +74,11 @@ class AgentContext(BaseModel):
|
|
|
66
74
|
description="Agent schema URI (e.g., 'rem-agents-query-agent')",
|
|
67
75
|
)
|
|
68
76
|
|
|
77
|
+
is_eval: bool = Field(
|
|
78
|
+
default=False,
|
|
79
|
+
description="Whether this is an evaluation session (set via X-Is-Eval header)",
|
|
80
|
+
)
|
|
81
|
+
|
|
69
82
|
model_config = {"populate_by_name": True}
|
|
70
83
|
|
|
71
84
|
@staticmethod
|
|
@@ -126,6 +139,7 @@ class AgentContext(BaseModel):
|
|
|
126
139
|
- X-Session-Id: Session identifier
|
|
127
140
|
- X-Model-Name: Model override
|
|
128
141
|
- X-Agent-Schema: Agent schema URI
|
|
142
|
+
- X-Is-Eval: Whether this is an evaluation session (true/false)
|
|
129
143
|
|
|
130
144
|
Args:
|
|
131
145
|
headers: Dictionary of HTTP headers (case-insensitive)
|
|
@@ -138,17 +152,23 @@ class AgentContext(BaseModel):
|
|
|
138
152
|
"X-User-Id": "user123",
|
|
139
153
|
"X-Tenant-Id": "acme-corp",
|
|
140
154
|
"X-Session-Id": "sess-456",
|
|
141
|
-
"X-Model-Name": "anthropic:claude-opus-4-20250514"
|
|
155
|
+
"X-Model-Name": "anthropic:claude-opus-4-20250514",
|
|
156
|
+
"X-Is-Eval": "true"
|
|
142
157
|
}
|
|
143
158
|
context = AgentContext.from_headers(headers)
|
|
144
159
|
"""
|
|
145
160
|
# Normalize header keys to lowercase for case-insensitive lookup
|
|
146
161
|
normalized = {k.lower(): v for k, v in headers.items()}
|
|
147
162
|
|
|
163
|
+
# Parse X-Is-Eval header (accepts "true", "1", "yes" as truthy)
|
|
164
|
+
is_eval_str = normalized.get("x-is-eval", "").lower()
|
|
165
|
+
is_eval = is_eval_str in ("true", "1", "yes")
|
|
166
|
+
|
|
148
167
|
return cls(
|
|
149
168
|
user_id=normalized.get("x-user-id"),
|
|
150
169
|
tenant_id=normalized.get("x-tenant-id", "default"),
|
|
151
170
|
session_id=normalized.get("x-session-id"),
|
|
152
171
|
default_model=normalized.get("x-model-name") or settings.llm.default_model,
|
|
153
172
|
agent_schema_uri=normalized.get("x-agent-schema"),
|
|
173
|
+
is_eval=is_eval,
|
|
154
174
|
)
|
rem/agentic/mcp/tool_wrapper.py
CHANGED
|
@@ -28,7 +28,12 @@ def create_pydantic_tool(func: Callable[..., Any]) -> Tool:
|
|
|
28
28
|
return Tool(func)
|
|
29
29
|
|
|
30
30
|
|
|
31
|
-
def create_mcp_tool_wrapper(
|
|
31
|
+
def create_mcp_tool_wrapper(
|
|
32
|
+
tool_name: str,
|
|
33
|
+
mcp_tool: Any,
|
|
34
|
+
user_id: str | None = None,
|
|
35
|
+
description_suffix: str | None = None,
|
|
36
|
+
) -> Tool:
|
|
32
37
|
"""
|
|
33
38
|
Create a Pydantic AI Tool from a FastMCP FunctionTool.
|
|
34
39
|
|
|
@@ -40,6 +45,8 @@ def create_mcp_tool_wrapper(tool_name: str, mcp_tool: Any, user_id: str | None =
|
|
|
40
45
|
tool_name: Name of the MCP tool
|
|
41
46
|
mcp_tool: The FastMCP FunctionTool object
|
|
42
47
|
user_id: Optional user_id to inject into tool calls
|
|
48
|
+
description_suffix: Optional text to append to the tool's docstring.
|
|
49
|
+
Used to add schema-specific context (e.g., default table for search_rem).
|
|
43
50
|
|
|
44
51
|
Returns:
|
|
45
52
|
A Pydantic AI Tool instance
|
|
@@ -52,7 +59,11 @@ def create_mcp_tool_wrapper(tool_name: str, mcp_tool: Any, user_id: str | None =
|
|
|
52
59
|
sig = inspect.signature(tool_func)
|
|
53
60
|
has_user_id = "user_id" in sig.parameters
|
|
54
61
|
|
|
55
|
-
#
|
|
62
|
+
# Build the docstring with optional suffix
|
|
63
|
+
base_doc = tool_func.__doc__ or ""
|
|
64
|
+
final_doc = base_doc + description_suffix if description_suffix else base_doc
|
|
65
|
+
|
|
66
|
+
# If we need to inject user_id or modify docstring, create a wrapper
|
|
56
67
|
# Otherwise, use the function directly for better signature preservation
|
|
57
68
|
if user_id and has_user_id:
|
|
58
69
|
async def wrapped_tool(**kwargs) -> Any:
|
|
@@ -69,12 +80,27 @@ def create_mcp_tool_wrapper(tool_name: str, mcp_tool: Any, user_id: str | None =
|
|
|
69
80
|
|
|
70
81
|
# Copy signature from original function for Pydantic AI inspection
|
|
71
82
|
wrapped_tool.__name__ = tool_name
|
|
72
|
-
wrapped_tool.__doc__ =
|
|
83
|
+
wrapped_tool.__doc__ = final_doc
|
|
73
84
|
wrapped_tool.__annotations__ = tool_func.__annotations__
|
|
74
85
|
wrapped_tool.__signature__ = sig # Important: preserve full signature
|
|
75
86
|
|
|
76
87
|
logger.debug(f"Creating MCP tool wrapper with user_id injection: {tool_name}")
|
|
77
88
|
return Tool(wrapped_tool)
|
|
89
|
+
elif description_suffix:
|
|
90
|
+
# Need to wrap just for docstring modification
|
|
91
|
+
async def wrapped_tool(**kwargs) -> Any:
|
|
92
|
+
"""Wrapper for docstring modification."""
|
|
93
|
+
valid_params = set(sig.parameters.keys())
|
|
94
|
+
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
|
|
95
|
+
return await tool_func(**filtered_kwargs)
|
|
96
|
+
|
|
97
|
+
wrapped_tool.__name__ = tool_name
|
|
98
|
+
wrapped_tool.__doc__ = final_doc
|
|
99
|
+
wrapped_tool.__annotations__ = tool_func.__annotations__
|
|
100
|
+
wrapped_tool.__signature__ = sig
|
|
101
|
+
|
|
102
|
+
logger.debug(f"Creating MCP tool wrapper with description suffix: {tool_name}")
|
|
103
|
+
return Tool(wrapped_tool)
|
|
78
104
|
else:
|
|
79
105
|
# No injection needed - use original function directly
|
|
80
106
|
logger.debug(f"Creating MCP tool wrapper (no injection): {tool_name}")
|
rem/agentic/otel/setup.py
CHANGED
|
@@ -591,6 +591,22 @@ async def create_agent(
|
|
|
591
591
|
|
|
592
592
|
set_agent_resource_attributes(agent_schema=agent_schema)
|
|
593
593
|
|
|
594
|
+
# Extract schema metadata for search_rem tool description suffix
|
|
595
|
+
# This allows entity schemas to add context-specific notes to the search_rem tool
|
|
596
|
+
search_rem_suffix = None
|
|
597
|
+
if metadata:
|
|
598
|
+
# Check for default_search_table in metadata (set by entity schemas)
|
|
599
|
+
extra = agent_schema.get("json_schema_extra", {}) if agent_schema else {}
|
|
600
|
+
default_table = extra.get("default_search_table")
|
|
601
|
+
has_embeddings = extra.get("has_embeddings", False)
|
|
602
|
+
|
|
603
|
+
if default_table:
|
|
604
|
+
# Build description suffix for search_rem
|
|
605
|
+
search_rem_suffix = f"\n\nFor this schema, use `search_rem` to query `{default_table}`. "
|
|
606
|
+
if has_embeddings:
|
|
607
|
+
search_rem_suffix += f"SEARCH works well on {default_table} (has embeddings). "
|
|
608
|
+
search_rem_suffix += f"Example: `SEARCH \"your query\" FROM {default_table} LIMIT 10`"
|
|
609
|
+
|
|
594
610
|
# Add tools from MCP server (in-process, no subprocess)
|
|
595
611
|
if mcp_server_configs:
|
|
596
612
|
for server_config in mcp_server_configs:
|
|
@@ -614,9 +630,17 @@ async def create_agent(
|
|
|
614
630
|
mcp_tools_dict = await mcp_server.get_tools()
|
|
615
631
|
|
|
616
632
|
for tool_name, tool_func in mcp_tools_dict.items():
|
|
617
|
-
|
|
633
|
+
# Add description suffix to search_rem tool if schema specifies a default table
|
|
634
|
+
tool_suffix = search_rem_suffix if tool_name == "search_rem" else None
|
|
635
|
+
|
|
636
|
+
wrapped_tool = create_mcp_tool_wrapper(
|
|
637
|
+
tool_name,
|
|
638
|
+
tool_func,
|
|
639
|
+
user_id=context.user_id if context else None,
|
|
640
|
+
description_suffix=tool_suffix,
|
|
641
|
+
)
|
|
618
642
|
tools.append(wrapped_tool)
|
|
619
|
-
logger.debug(f"Loaded MCP tool: {tool_name}")
|
|
643
|
+
logger.debug(f"Loaded MCP tool: {tool_name}" + (" (with schema suffix)" if tool_suffix else ""))
|
|
620
644
|
|
|
621
645
|
logger.info(f"Loaded {len(mcp_tools_dict)} tools from MCP server: {server_id} (in-process)")
|
|
622
646
|
|
rem/api/main.py
CHANGED
|
@@ -376,10 +376,13 @@ def create_app() -> FastAPI:
|
|
|
376
376
|
|
|
377
377
|
app.include_router(chat_router)
|
|
378
378
|
app.include_router(models_router)
|
|
379
|
+
# shared_sessions_router MUST be before messages_router
|
|
380
|
+
# because messages_router has /sessions/{session_id} which would match
|
|
381
|
+
# before the more specific /sessions/shared-with-me routes
|
|
382
|
+
app.include_router(shared_sessions_router)
|
|
379
383
|
app.include_router(messages_router)
|
|
380
384
|
app.include_router(feedback_router)
|
|
381
385
|
app.include_router(admin_router)
|
|
382
|
-
app.include_router(shared_sessions_router)
|
|
383
386
|
app.include_router(query_router)
|
|
384
387
|
|
|
385
388
|
# Register auth router (if enabled)
|
rem/api/mcp_router/server.py
CHANGED
|
@@ -127,10 +127,12 @@ def create_mcp_server(is_local: bool = False) -> FastMCP:
|
|
|
127
127
|
"AVAILABLE TOOLS\n"
|
|
128
128
|
"═══════════════════════════════════════════════════════════════════════════\n"
|
|
129
129
|
"\n"
|
|
130
|
-
"•
|
|
131
|
-
"•
|
|
130
|
+
"• search_rem - Execute REM queries (LOOKUP, FUZZY, SEARCH, SQL, TRAVERSE)\n"
|
|
131
|
+
"• ask_rem_agent - Natural language to REM query conversion\n"
|
|
132
132
|
" - plan_mode=True: Hints agent to use TRAVERSE with depth=0 for edge analysis\n"
|
|
133
|
-
"•
|
|
133
|
+
"• ingest_into_rem - Ingest files from local paths (local server only), s3://, or https://\n"
|
|
134
|
+
"• list_schema - List all database schemas (tables) with row counts\n"
|
|
135
|
+
"• get_schema - Get detailed schema for a specific table (columns, types, indexes)\n"
|
|
134
136
|
"\n"
|
|
135
137
|
"═══════════════════════════════════════════════════════════════════════════\n"
|
|
136
138
|
"AVAILABLE RESOURCES (Read-Only)\n"
|
|
@@ -175,7 +177,9 @@ def create_mcp_server(is_local: bool = False) -> FastMCP:
|
|
|
175
177
|
# Register REM tools
|
|
176
178
|
from .tools import (
|
|
177
179
|
ask_rem_agent,
|
|
180
|
+
get_schema,
|
|
178
181
|
ingest_into_rem,
|
|
182
|
+
list_schema,
|
|
179
183
|
read_resource,
|
|
180
184
|
register_metadata,
|
|
181
185
|
search_rem,
|
|
@@ -185,6 +189,8 @@ def create_mcp_server(is_local: bool = False) -> FastMCP:
|
|
|
185
189
|
mcp.tool()(ask_rem_agent)
|
|
186
190
|
mcp.tool()(read_resource)
|
|
187
191
|
mcp.tool()(register_metadata)
|
|
192
|
+
mcp.tool()(list_schema)
|
|
193
|
+
mcp.tool()(get_schema)
|
|
188
194
|
|
|
189
195
|
# File ingestion tool (with local path support for local servers)
|
|
190
196
|
# Wrap to inject is_local parameter
|
rem/api/mcp_router/tools.py
CHANGED
|
@@ -15,6 +15,9 @@ Available Tools:
|
|
|
15
15
|
- ask_rem_agent: Natural language to REM query conversion via agent
|
|
16
16
|
- ingest_into_rem: Full file ingestion pipeline (read + store + parse + chunk)
|
|
17
17
|
- read_resource: Access MCP resources (for Claude Desktop compatibility)
|
|
18
|
+
- register_metadata: Register response metadata for SSE MetadataEvent
|
|
19
|
+
- list_schema: List all schemas (tables, agents) in the database with row counts
|
|
20
|
+
- get_schema: Get detailed schema for a table (columns, types, indexes)
|
|
18
21
|
"""
|
|
19
22
|
|
|
20
23
|
from functools import wraps
|
|
@@ -603,7 +606,9 @@ async def register_metadata(
|
|
|
603
606
|
references: list[str] | None = None,
|
|
604
607
|
sources: list[str] | None = None,
|
|
605
608
|
flags: list[str] | None = None,
|
|
606
|
-
#
|
|
609
|
+
# Session naming
|
|
610
|
+
session_name: str | None = None,
|
|
611
|
+
# Risk assessment fields (used by specialized agents)
|
|
607
612
|
risk_level: str | None = None,
|
|
608
613
|
risk_score: int | None = None,
|
|
609
614
|
risk_reasoning: str | None = None,
|
|
@@ -636,6 +641,11 @@ async def register_metadata(
|
|
|
636
641
|
flags: Optional flags for the response (e.g., "needs_review",
|
|
637
642
|
"uncertain", "incomplete", "crisis_alert").
|
|
638
643
|
|
|
644
|
+
session_name: Short 1-3 phrase name describing the session topic.
|
|
645
|
+
Used by the UI to label conversations in the sidebar.
|
|
646
|
+
Examples: "Prescription Drug Questions", "AWS Setup Help",
|
|
647
|
+
"Python Code Review", "Travel Planning".
|
|
648
|
+
|
|
639
649
|
risk_level: Risk level indicator (e.g., "green", "orange", "red").
|
|
640
650
|
Used by mental health agents for C-SSRS style assessment.
|
|
641
651
|
risk_score: Numeric risk score (e.g., 0-6 for C-SSRS).
|
|
@@ -660,7 +670,7 @@ async def register_metadata(
|
|
|
660
670
|
sources=["REM database lookup"]
|
|
661
671
|
)
|
|
662
672
|
|
|
663
|
-
#
|
|
673
|
+
# Risk assessment example
|
|
664
674
|
register_metadata(
|
|
665
675
|
confidence=0.9,
|
|
666
676
|
risk_level="green",
|
|
@@ -703,6 +713,10 @@ async def register_metadata(
|
|
|
703
713
|
"flags": flags,
|
|
704
714
|
}
|
|
705
715
|
|
|
716
|
+
# Add session name if provided
|
|
717
|
+
if session_name is not None:
|
|
718
|
+
result["session_name"] = session_name
|
|
719
|
+
|
|
706
720
|
# Add risk assessment fields if provided
|
|
707
721
|
if risk_level is not None:
|
|
708
722
|
result["risk_level"] = risk_level
|
|
@@ -718,3 +732,311 @@ async def register_metadata(
|
|
|
718
732
|
result["extra"] = extra
|
|
719
733
|
|
|
720
734
|
return result
|
|
735
|
+
|
|
736
|
+
|
|
737
|
+
@mcp_tool_error_handler
|
|
738
|
+
async def list_schema(
|
|
739
|
+
include_system: bool = False,
|
|
740
|
+
user_id: str | None = None,
|
|
741
|
+
) -> dict[str, Any]:
|
|
742
|
+
"""
|
|
743
|
+
List all schemas (tables) in the REM database.
|
|
744
|
+
|
|
745
|
+
Returns metadata about all available tables including their names,
|
|
746
|
+
row counts, and descriptions. Use this to discover what data is
|
|
747
|
+
available before constructing queries.
|
|
748
|
+
|
|
749
|
+
Args:
|
|
750
|
+
include_system: If True, include PostgreSQL system tables (pg_*, information_schema).
|
|
751
|
+
Default False shows only REM application tables.
|
|
752
|
+
user_id: Optional user identifier (defaults to authenticated user or "default")
|
|
753
|
+
|
|
754
|
+
Returns:
|
|
755
|
+
Dict with:
|
|
756
|
+
- status: "success" or "error"
|
|
757
|
+
- tables: List of table metadata dicts with:
|
|
758
|
+
- name: Table name
|
|
759
|
+
- schema: Schema name (usually "public")
|
|
760
|
+
- estimated_rows: Approximate row count
|
|
761
|
+
- description: Table comment if available
|
|
762
|
+
|
|
763
|
+
Examples:
|
|
764
|
+
# List all REM schemas
|
|
765
|
+
list_schema()
|
|
766
|
+
|
|
767
|
+
# Include system tables
|
|
768
|
+
list_schema(include_system=True)
|
|
769
|
+
"""
|
|
770
|
+
rem_service = await get_rem_service()
|
|
771
|
+
user_id = AgentContext.get_user_id_or_default(user_id, source="list_schema")
|
|
772
|
+
|
|
773
|
+
# Query information_schema for tables
|
|
774
|
+
schema_filter = ""
|
|
775
|
+
if not include_system:
|
|
776
|
+
schema_filter = """
|
|
777
|
+
AND table_schema = 'public'
|
|
778
|
+
AND table_name NOT LIKE 'pg_%'
|
|
779
|
+
AND table_name NOT LIKE '_pg_%'
|
|
780
|
+
"""
|
|
781
|
+
|
|
782
|
+
query = f"""
|
|
783
|
+
SELECT
|
|
784
|
+
t.table_schema,
|
|
785
|
+
t.table_name,
|
|
786
|
+
pg_catalog.obj_description(
|
|
787
|
+
(quote_ident(t.table_schema) || '.' || quote_ident(t.table_name))::regclass,
|
|
788
|
+
'pg_class'
|
|
789
|
+
) as description,
|
|
790
|
+
(
|
|
791
|
+
SELECT reltuples::bigint
|
|
792
|
+
FROM pg_class c
|
|
793
|
+
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
794
|
+
WHERE c.relname = t.table_name
|
|
795
|
+
AND n.nspname = t.table_schema
|
|
796
|
+
) as estimated_rows
|
|
797
|
+
FROM information_schema.tables t
|
|
798
|
+
WHERE t.table_type = 'BASE TABLE'
|
|
799
|
+
{schema_filter}
|
|
800
|
+
ORDER BY t.table_schema, t.table_name
|
|
801
|
+
"""
|
|
802
|
+
|
|
803
|
+
# Access postgres service directly from cache
|
|
804
|
+
postgres_service = _service_cache.get("postgres")
|
|
805
|
+
if not postgres_service:
|
|
806
|
+
postgres_service = rem_service._postgres
|
|
807
|
+
|
|
808
|
+
rows = await postgres_service.fetch(query)
|
|
809
|
+
|
|
810
|
+
tables = []
|
|
811
|
+
for row in rows:
|
|
812
|
+
tables.append({
|
|
813
|
+
"name": row["table_name"],
|
|
814
|
+
"schema": row["table_schema"],
|
|
815
|
+
"estimated_rows": int(row["estimated_rows"]) if row["estimated_rows"] else 0,
|
|
816
|
+
"description": row["description"],
|
|
817
|
+
})
|
|
818
|
+
|
|
819
|
+
logger.info(f"Listed {len(tables)} schemas for user {user_id}")
|
|
820
|
+
|
|
821
|
+
return {
|
|
822
|
+
"tables": tables,
|
|
823
|
+
"count": len(tables),
|
|
824
|
+
}
|
|
825
|
+
|
|
826
|
+
|
|
827
|
+
@mcp_tool_error_handler
|
|
828
|
+
async def get_schema(
|
|
829
|
+
table_name: str,
|
|
830
|
+
include_indexes: bool = True,
|
|
831
|
+
include_constraints: bool = True,
|
|
832
|
+
columns: list[str] | None = None,
|
|
833
|
+
user_id: str | None = None,
|
|
834
|
+
) -> dict[str, Any]:
|
|
835
|
+
"""
|
|
836
|
+
Get detailed schema information for a specific table.
|
|
837
|
+
|
|
838
|
+
Returns column definitions, data types, constraints, and indexes.
|
|
839
|
+
Use this to understand table structure before writing SQL queries.
|
|
840
|
+
|
|
841
|
+
Args:
|
|
842
|
+
table_name: Name of the table to inspect (e.g., "resources", "moments")
|
|
843
|
+
include_indexes: Include index information (default True)
|
|
844
|
+
include_constraints: Include constraint information (default True)
|
|
845
|
+
columns: Optional list of specific columns to return. If None, returns all columns.
|
|
846
|
+
user_id: Optional user identifier (defaults to authenticated user or "default")
|
|
847
|
+
|
|
848
|
+
Returns:
|
|
849
|
+
Dict with:
|
|
850
|
+
- status: "success" or "error"
|
|
851
|
+
- table_name: Name of the table
|
|
852
|
+
- columns: List of column definitions with:
|
|
853
|
+
- name: Column name
|
|
854
|
+
- type: PostgreSQL data type
|
|
855
|
+
- nullable: Whether NULL is allowed
|
|
856
|
+
- default: Default value if any
|
|
857
|
+
- description: Column comment if available
|
|
858
|
+
- indexes: List of indexes (if include_indexes=True)
|
|
859
|
+
- constraints: List of constraints (if include_constraints=True)
|
|
860
|
+
- primary_key: Primary key column(s)
|
|
861
|
+
|
|
862
|
+
Examples:
|
|
863
|
+
# Get full schema for resources table
|
|
864
|
+
get_schema(table_name="resources")
|
|
865
|
+
|
|
866
|
+
# Get only specific columns
|
|
867
|
+
get_schema(
|
|
868
|
+
table_name="resources",
|
|
869
|
+
columns=["id", "name", "created_at"]
|
|
870
|
+
)
|
|
871
|
+
|
|
872
|
+
# Get schema without indexes
|
|
873
|
+
get_schema(
|
|
874
|
+
table_name="moments",
|
|
875
|
+
include_indexes=False
|
|
876
|
+
)
|
|
877
|
+
"""
|
|
878
|
+
rem_service = await get_rem_service()
|
|
879
|
+
user_id = AgentContext.get_user_id_or_default(user_id, source="get_schema")
|
|
880
|
+
|
|
881
|
+
# Access postgres service
|
|
882
|
+
postgres_service = _service_cache.get("postgres")
|
|
883
|
+
if not postgres_service:
|
|
884
|
+
postgres_service = rem_service._postgres
|
|
885
|
+
|
|
886
|
+
# Verify table exists
|
|
887
|
+
exists_query = """
|
|
888
|
+
SELECT EXISTS (
|
|
889
|
+
SELECT 1 FROM information_schema.tables
|
|
890
|
+
WHERE table_schema = 'public' AND table_name = $1
|
|
891
|
+
)
|
|
892
|
+
"""
|
|
893
|
+
exists = await postgres_service.fetchval(exists_query, table_name)
|
|
894
|
+
if not exists:
|
|
895
|
+
return {
|
|
896
|
+
"status": "error",
|
|
897
|
+
"error": f"Table '{table_name}' not found in public schema",
|
|
898
|
+
}
|
|
899
|
+
|
|
900
|
+
# Get columns
|
|
901
|
+
columns_filter = ""
|
|
902
|
+
if columns:
|
|
903
|
+
placeholders = ", ".join(f"${i+2}" for i in range(len(columns)))
|
|
904
|
+
columns_filter = f"AND column_name IN ({placeholders})"
|
|
905
|
+
|
|
906
|
+
columns_query = f"""
|
|
907
|
+
SELECT
|
|
908
|
+
c.column_name,
|
|
909
|
+
c.data_type,
|
|
910
|
+
c.udt_name,
|
|
911
|
+
c.is_nullable,
|
|
912
|
+
c.column_default,
|
|
913
|
+
c.character_maximum_length,
|
|
914
|
+
c.numeric_precision,
|
|
915
|
+
pg_catalog.col_description(
|
|
916
|
+
(quote_ident(c.table_schema) || '.' || quote_ident(c.table_name))::regclass,
|
|
917
|
+
c.ordinal_position
|
|
918
|
+
) as description
|
|
919
|
+
FROM information_schema.columns c
|
|
920
|
+
WHERE c.table_schema = 'public'
|
|
921
|
+
AND c.table_name = $1
|
|
922
|
+
{columns_filter}
|
|
923
|
+
ORDER BY c.ordinal_position
|
|
924
|
+
"""
|
|
925
|
+
|
|
926
|
+
params = [table_name]
|
|
927
|
+
if columns:
|
|
928
|
+
params.extend(columns)
|
|
929
|
+
|
|
930
|
+
column_rows = await postgres_service.fetch(columns_query, *params)
|
|
931
|
+
|
|
932
|
+
column_defs = []
|
|
933
|
+
for row in column_rows:
|
|
934
|
+
# Build a more readable type string
|
|
935
|
+
data_type = row["data_type"]
|
|
936
|
+
if row["character_maximum_length"]:
|
|
937
|
+
data_type = f"{data_type}({row['character_maximum_length']})"
|
|
938
|
+
elif row["udt_name"] in ("int4", "int8", "float4", "float8"):
|
|
939
|
+
# Use common type names
|
|
940
|
+
type_map = {"int4": "integer", "int8": "bigint", "float4": "real", "float8": "double precision"}
|
|
941
|
+
data_type = type_map.get(row["udt_name"], data_type)
|
|
942
|
+
elif row["udt_name"] == "vector":
|
|
943
|
+
data_type = "vector"
|
|
944
|
+
|
|
945
|
+
column_defs.append({
|
|
946
|
+
"name": row["column_name"],
|
|
947
|
+
"type": data_type,
|
|
948
|
+
"nullable": row["is_nullable"] == "YES",
|
|
949
|
+
"default": row["column_default"],
|
|
950
|
+
"description": row["description"],
|
|
951
|
+
})
|
|
952
|
+
|
|
953
|
+
result = {
|
|
954
|
+
"table_name": table_name,
|
|
955
|
+
"columns": column_defs,
|
|
956
|
+
"column_count": len(column_defs),
|
|
957
|
+
}
|
|
958
|
+
|
|
959
|
+
# Get primary key
|
|
960
|
+
pk_query = """
|
|
961
|
+
SELECT a.attname as column_name
|
|
962
|
+
FROM pg_index i
|
|
963
|
+
JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
|
|
964
|
+
WHERE i.indrelid = $1::regclass
|
|
965
|
+
AND i.indisprimary
|
|
966
|
+
ORDER BY array_position(i.indkey, a.attnum)
|
|
967
|
+
"""
|
|
968
|
+
pk_rows = await postgres_service.fetch(pk_query, table_name)
|
|
969
|
+
result["primary_key"] = [row["column_name"] for row in pk_rows]
|
|
970
|
+
|
|
971
|
+
# Get indexes
|
|
972
|
+
if include_indexes:
|
|
973
|
+
indexes_query = """
|
|
974
|
+
SELECT
|
|
975
|
+
i.relname as index_name,
|
|
976
|
+
am.amname as index_type,
|
|
977
|
+
ix.indisunique as is_unique,
|
|
978
|
+
ix.indisprimary as is_primary,
|
|
979
|
+
array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
|
|
980
|
+
FROM pg_index ix
|
|
981
|
+
JOIN pg_class i ON i.oid = ix.indexrelid
|
|
982
|
+
JOIN pg_class t ON t.oid = ix.indrelid
|
|
983
|
+
JOIN pg_am am ON am.oid = i.relam
|
|
984
|
+
JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
|
|
985
|
+
WHERE t.relname = $1
|
|
986
|
+
GROUP BY i.relname, am.amname, ix.indisunique, ix.indisprimary
|
|
987
|
+
ORDER BY i.relname
|
|
988
|
+
"""
|
|
989
|
+
index_rows = await postgres_service.fetch(indexes_query, table_name)
|
|
990
|
+
result["indexes"] = [
|
|
991
|
+
{
|
|
992
|
+
"name": row["index_name"],
|
|
993
|
+
"type": row["index_type"],
|
|
994
|
+
"unique": row["is_unique"],
|
|
995
|
+
"primary": row["is_primary"],
|
|
996
|
+
"columns": row["columns"],
|
|
997
|
+
}
|
|
998
|
+
for row in index_rows
|
|
999
|
+
]
|
|
1000
|
+
|
|
1001
|
+
# Get constraints
|
|
1002
|
+
if include_constraints:
|
|
1003
|
+
constraints_query = """
|
|
1004
|
+
SELECT
|
|
1005
|
+
con.conname as constraint_name,
|
|
1006
|
+
con.contype as constraint_type,
|
|
1007
|
+
array_agg(a.attname ORDER BY array_position(con.conkey, a.attnum)) as columns,
|
|
1008
|
+
pg_get_constraintdef(con.oid) as definition
|
|
1009
|
+
FROM pg_constraint con
|
|
1010
|
+
JOIN pg_class t ON t.oid = con.conrelid
|
|
1011
|
+
JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(con.conkey)
|
|
1012
|
+
WHERE t.relname = $1
|
|
1013
|
+
GROUP BY con.conname, con.contype, con.oid
|
|
1014
|
+
ORDER BY con.contype, con.conname
|
|
1015
|
+
"""
|
|
1016
|
+
constraint_rows = await postgres_service.fetch(constraints_query, table_name)
|
|
1017
|
+
|
|
1018
|
+
# Map constraint types to readable names
|
|
1019
|
+
type_map = {
|
|
1020
|
+
"p": "PRIMARY KEY",
|
|
1021
|
+
"u": "UNIQUE",
|
|
1022
|
+
"f": "FOREIGN KEY",
|
|
1023
|
+
"c": "CHECK",
|
|
1024
|
+
"x": "EXCLUSION",
|
|
1025
|
+
}
|
|
1026
|
+
|
|
1027
|
+
result["constraints"] = []
|
|
1028
|
+
for row in constraint_rows:
|
|
1029
|
+
# contype is returned as bytes (char type), decode it
|
|
1030
|
+
con_type = row["constraint_type"]
|
|
1031
|
+
if isinstance(con_type, bytes):
|
|
1032
|
+
con_type = con_type.decode("utf-8")
|
|
1033
|
+
result["constraints"].append({
|
|
1034
|
+
"name": row["constraint_name"],
|
|
1035
|
+
"type": type_map.get(con_type, con_type),
|
|
1036
|
+
"columns": row["columns"],
|
|
1037
|
+
"definition": row["definition"],
|
|
1038
|
+
})
|
|
1039
|
+
|
|
1040
|
+
logger.info(f"Retrieved schema for table '{table_name}' with {len(column_defs)} columns")
|
|
1041
|
+
|
|
1042
|
+
return result
|