From 44cad7249d814f4d86efe06f913dcf9844365869 Mon Sep 17 00:00:00 2001 From: yuvalkh Date: Wed, 5 Aug 2026 14:52:59 +0300 Subject: [PATCH 1/2] added option to choose connection by name from jeen metadata --- agent/src/agent/mcp_server.py | 2 + agent/src/agent/nodes/detect_ambiguity.py | 6 +- agent/src/agent/nodes/schema_explorer.py | 3 +- agent/src/agent/state.py | 1 + agent/src/agent/utils/jeen_metadata_client.py | 15 +++-- backend/app/config.py | 6 ++ backend/app/routers/agent.py | 37 ++++++++++++ frontend/src/api/agent.ts | 5 ++ frontend/src/pages/AgentTestingPage.tsx | 59 ++++++++++++------- 9 files changed, 104 insertions(+), 30 deletions(-) diff --git a/agent/src/agent/mcp_server.py b/agent/src/agent/mcp_server.py index e937ff9..fcdf078 100644 --- a/agent/src/agent/mcp_server.py +++ b/agent/src/agent/mcp_server.py @@ -24,6 +24,7 @@ async def chat_with_agent( active_skills: list[str] | None = None, execution_mode: str | None = None, hitl_enabled: bool = True, + connection_id: int | None = None, ) -> str: """Run the Text2SQL agent to answer database queries. @@ -139,6 +140,7 @@ async def chat_with_agent( "active_skills": active_skills, "execution_mode": execution_mode, "non_interactive": not hitl_enabled, + "connection_id": connection_id, }, config=config, ) diff --git a/agent/src/agent/nodes/detect_ambiguity.py b/agent/src/agent/nodes/detect_ambiguity.py index 3243372..2c01955 100644 --- a/agent/src/agent/nodes/detect_ambiguity.py +++ b/agent/src/agent/nodes/detect_ambiguity.py @@ -59,8 +59,8 @@ def _resolve_ambiguity_type(parsed: dict) -> str: ambiguity_type = parsed.get("ambiguity_type") clarifying: str | None = parsed.get("clarifying_questions") - # If the LLM classified it as unanswerable but generated a clarifying question, it MUST be ambiguous. - if ambiguity_type == "unanswerable" and clarifying is not None and clarifying.strip(): + # If the LLM generated a clarifying question, it MUST be ambiguous. + if clarifying is not None and isinstance(clarifying, str) and clarifying.strip(): return "ambiguous" if ambiguity_type in ["clear", "ambiguous", "unanswerable"]: @@ -70,10 +70,10 @@ def _resolve_ambiguity_type(parsed: dict) -> str: class AmbiguityResult(BaseModel): + reason: str = Field(default="", description="Explanation of the ambiguity detection decision. Always think through your reasoning here first.") ambiguity_type: Literal["clear", "ambiguous", "unanswerable"] = Field( description="The determined state of the query: 'clear' (proceed), 'ambiguous' (needs clarification), or 'unanswerable' (impossible)." ) - reason: str = Field(default="", description="Explanation of the ambiguity detection decision") clarifying_questions: str | None = Field(default=None, description="Questions to ask the user if ambiguous. Null if clear or unanswerable.") diff --git a/agent/src/agent/nodes/schema_explorer.py b/agent/src/agent/nodes/schema_explorer.py index ed3d976..36dd5a6 100644 --- a/agent/src/agent/nodes/schema_explorer.py +++ b/agent/src/agent/nodes/schema_explorer.py @@ -32,7 +32,8 @@ async def schema_explorer_node(state: AgentState, config: RunnableConfig | None logger.info("Fetching full catalog prompt from Jeen MCP.") try: - catalog_prompt = await _jeen.get_catalog_prompt() + connection_id = state.get("connection_id") + catalog_prompt = await _jeen.get_catalog_prompt(connection_id=connection_id) if not catalog_prompt: raise ValueError("Received empty catalog prompt from Jeen.") except Exception as exc: diff --git a/agent/src/agent/state.py b/agent/src/agent/state.py index cc49c34..0b71325 100644 --- a/agent/src/agent/state.py +++ b/agent/src/agent/state.py @@ -25,6 +25,7 @@ class AgentState(TypedDict): active_extractors: list[dict[str, str]] | None active_skills: list[str] | None loaded_skills: list[dict] | None + connection_id: int | None last_error: str | None esca_write_failed: bool | None inline_result_rows: list[list[Any]] | None diff --git a/agent/src/agent/utils/jeen_metadata_client.py b/agent/src/agent/utils/jeen_metadata_client.py index 8d7cd1d..251cb53 100644 --- a/agent/src/agent/utils/jeen_metadata_client.py +++ b/agent/src/agent/utils/jeen_metadata_client.py @@ -155,17 +155,18 @@ async def _call(self, tool_name: str, arguments: dict[str, Any]) -> Any: # Full DB Schema (all columns) # ------------------------------------------------------------------ - async def get_catalog_prompt(self) -> str: + async def get_catalog_prompt(self, connection_id: int | None = None) -> str: """ Fetch the entire catalog context prompt for the connection using the MCP `get_catalog_prompt` tool. This returns a large markdown string describing all tables, columns, relationships, and business terms. """ try: + cid = connection_id if connection_id is not None else self._connection_id payload = await self._call( "get_catalog_prompt", { - "connection_id": self._connection_id, + "connection_id": cid, }, ) # The MCP tool returns { "content": [{ "type": "text", "text": "..." }] } @@ -192,7 +193,7 @@ async def get_catalog_prompt(self) -> str: # Table profile (columns + stats) # ------------------------------------------------------------------ - async def get_table_profile(self, table_name: str) -> dict[str, Any] | None: + async def get_table_profile(self, table_name: str, connection_id: int | None = None) -> dict[str, Any] | None: """ Fetch the latest stored column stats for *table_name* from jeen-metadata. @@ -213,10 +214,11 @@ async def get_table_profile(self, table_name: str) -> dict[str, Any] | None: } """ try: + cid = connection_id if connection_id is not None else self._connection_id payload = await self._call( "get_table_profile", { - "connection_id": self._connection_id, + "connection_id": cid, "table_name": table_name, }, ) @@ -277,15 +279,16 @@ async def get_table_profile(self, table_name: str) -> dict[str, Any] | None: # Full table listing (fallback when search returns nothing) # ------------------------------------------------------------------ - async def list_tables_rich(self) -> list[dict[str, Any]]: + async def list_tables_rich(self, connection_id: int | None = None) -> list[dict[str, Any]]: """ Return ALL tables for the configured connection via ``list_tables_rich``. Used as a fallback when the search tool returns no results. """ try: + cid = connection_id if connection_id is not None else self._connection_id rows = await self._call( "list_tables_rich", - {"connection_id": self._connection_id}, + {"connection_id": cid}, ) if not isinstance(rows, list): rows = [] diff --git a/backend/app/config.py b/backend/app/config.py index eec6162..d53e60c 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -32,6 +32,12 @@ class Settings(BaseSettings): EVALUATION_SERVICE_URL: str = "http://localhost:8001" OPENMETADATA_TOKEN: str = "" + # Jeen MCP + JEEN_METADATA_MCP_URL: str = "http://schema-modeler.dev161.internal/api/mcp" + JEEN_METADATA_MCP_KEY: str = ( + "mcp_f885337e381366db5edc22093415450e38f71e997e96dc708fea69bde9529ab9" + ) + APP_ENV: str = "development" OPENMETADATA_URL: str = "http://localhost:8585" OPENMETADATA_ADMIN_EMAIL: str = "admin@open-metadata.org" diff --git a/backend/app/routers/agent.py b/backend/app/routers/agent.py index 6994420..1e51127 100644 --- a/backend/app/routers/agent.py +++ b/backend/app/routers/agent.py @@ -48,6 +48,7 @@ class ChatRequest(BaseModel): active_skills: list[str] | None = None execution_mode: str | None = None hitl_enabled: bool = True + connection_id: int | None = None class SuggestFixesRequest(BaseModel): @@ -83,6 +84,24 @@ async def _get_mcp_client(): yield session +@asynccontextmanager +async def _get_jeen_mcp_client(): + url = settings.JEEN_METADATA_MCP_URL + headers = ( + {"Authorization": f"Bearer {settings.JEEN_METADATA_MCP_KEY}"} + if settings.JEEN_METADATA_MCP_KEY + else {} + ) + async with streamablehttp_client(url, headers=headers) as ( + read_stream, + write_stream, + _, + ): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + yield session + + async def _call_agent_mcp(tool_arguments: dict) -> dict: """ Connects to the agent MCP server over Streamable HTTP, initializes the session, @@ -189,6 +208,7 @@ async def chat(request: ChatRequest) -> ChatResponse: "allowed_statuses": request.allowed_statuses, "extractors": request.extractors, "hitl_enabled": request.hitl_enabled, + "connection_id": request.connection_id, } result = await _call_agent_mcp(tool_arguments) @@ -323,3 +343,20 @@ async def suggest_fixes(req: SuggestFixesRequest): except Exception as e: logger.error(f"Suggest fixes error: {e}") return [] + + +@router.get("/connections") +async def list_connections(): + """Fetch list of connections from the Jeen MCP.""" + try: + async with _get_jeen_mcp_client() as session: + result = await session.call_tool( + "list_connections", + arguments={}, + read_timeout_seconds=timedelta(seconds=60.0), + ) + content = result.content[0].text + return json.loads(content) + except Exception as e: + logger.error(f"List connections error: {e}") + return {"connections": []} diff --git a/frontend/src/api/agent.ts b/frontend/src/api/agent.ts index cd37bee..b910f75 100644 --- a/frontend/src/api/agent.ts +++ b/frontend/src/api/agent.ts @@ -32,6 +32,7 @@ export interface ChatRequest { allowed_statuses?: string[]; extractors?: string[]; hitl_enabled?: boolean; + connection_id?: number; } export interface ChatResponse { @@ -54,4 +55,8 @@ export const agentApi = { const response = await api.post('/suggest_fixes', { thread_id: threadId, category }); return response.data; }, + listConnections: async (): Promise => { + const response = await api.get('/connections'); + return response.data; + }, }; diff --git a/frontend/src/pages/AgentTestingPage.tsx b/frontend/src/pages/AgentTestingPage.tsx index dfbf9b9..b1813b9 100644 --- a/frontend/src/pages/AgentTestingPage.tsx +++ b/frontend/src/pages/AgentTestingPage.tsx @@ -354,13 +354,17 @@ const AgentTestingHeader = memo( ({ hitlEnabled, setHitlEnabled, - allowedStatuses, - setAllowedStatuses, + connections, + selectedConnection, + setSelectedConnection, + loadingConnections, }: { hitlEnabled: boolean; setHitlEnabled: (v: boolean) => void; - allowedStatuses: string[]; - setAllowedStatuses: (v: string[]) => void; + connections: any[]; + selectedConnection: number | undefined; + setSelectedConnection: (v: number | undefined) => void; + loadingConnections: boolean; }) => (
@@ -378,20 +382,17 @@ const AgentTestingHeader = memo(
- Table Status + Catalog