added option to choose connection by name from jeen metadata

agent/src/agent/mcp_server.py

@@ -15,6 +15,7 @@
 @mcp.tool()
 @observe()
 async def chat_with_agent(
+    connection_id: int,
     query: str | None = None,
     thread_id: str | None = None,
     resume_value: str | dict | None = None,
@@ -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,
         )

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.")
 
 

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:

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
     last_error: str | None
     esca_write_failed: bool | None
     inline_result_rows: list[list[Any]] | None

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 = []

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"

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
 
 
 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": []}

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<string[]>('/suggest_fixes', { thread_id: threadId, category });
     return response.data;
   },
+  listConnections: async (): Promise<any> => {
+    const response = await api.get<any>('/connections');
+    return response.data;
+  },
 };

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;
   }) => (
     <div className={styles.header}>
       <div>
@@ -378,20 +382,17 @@ const AgentTestingHeader = memo(
           <Switch checked={hitlEnabled} onChange={setHitlEnabled} />
         </div>
         <div className={styles.controlItem}>
-          <span>Table Status</span>
+          <span>Catalog</span>
           <Select
-            mode="multiple"
-            value={allowedStatuses}
-            onChange={setAllowedStatuses}
+            value={selectedConnection}
+            onChange={setSelectedConnection}
             style={{ minWidth: 200 }}
-            placeholder="Select allowed statuses"
-            options={[
-              { value: 'production', label: 'Production' },
-              { value: 'verified', label: 'Verified' },
-              { value: 'sandbox', label: 'Sandbox' },
-              { value: 'draft', label: 'Draft' },
-              { value: 'degraded', label: 'Degraded' },
-            ]}
+            placeholder="Select connection"
+            loading={loadingConnections}
+            options={connections.map((c: any) => ({
+              value: c.connection_id,
+              label: c.name,
+            }))}
           />
         </div>
       </div>
@@ -409,12 +410,14 @@ const AgentChatInput = ({
   onSubmit,
   disabled,
   loading,
+  submitDisabled,
 }: {
   query: string;
   setQuery: (q: string) => void;
   onSubmit: () => void;
   disabled: boolean;
   loading: boolean;
+  submitDisabled?: boolean;
 }) => (
   <div className={`${styles.glassCard} ${styles.animateIn}`}>
     <Space.Compact className={styles.chatInputWrapper}>
@@ -423,14 +426,16 @@ const AgentChatInput = ({
         placeholder="Ask the agent to query a table..."
         value={query}
         onChange={(e) => setQuery(e.target.value)}
-        onPressEnter={onSubmit}
+        onPressEnter={(e) => {
+          if (!disabled && !submitDisabled) onSubmit();
+        }}
         disabled={disabled}
       />
       <Button
         className={styles.primaryButton}
         onClick={onSubmit}
         loading={loading}
-        disabled={disabled}
+        disabled={disabled || submitDisabled}
       >
         {!loading && <Play size={18} />}
         Run Agent
@@ -892,7 +897,23 @@ const AgentResultDisplay = ({
 export function AgentTestingPage() {
   const [query, setQuery] = useState('');
   const [hitlEnabled, setHitlEnabled] = useState(true);
-  const [allowedStatuses, setAllowedStatuses] = useState<string[]>(['production']);
+  const [connections, setConnections] = useState<any[]>([]);
+  const [selectedConnection, setSelectedConnection] = useState<number | undefined>(undefined);
+  const [loadingConnections, setLoadingConnections] = useState(false);
+
+  useEffect(() => {
+    setLoadingConnections(true);
+    agentApi
+      .listConnections()
+      .then((res) => {
+        setConnections(res.connections || []);
+        setLoadingConnections(false);
+      })
+      .catch((err) => {
+        console.error('Failed to fetch connections:', err);
+        setLoadingConnections(false);
+      });
+  }, []);
   const [threadId, setThreadId] = useState<string | null>(null);
   const [traceId, setTraceId] = useState<string | null>(null);
   const [chatResponse, setChatResponse] = useState<ChatResponse | null>(null);
@@ -1004,7 +1025,7 @@ export function AgentTestingPage() {
   const isInputDisabled = chatMutation.isPending || chatResponse?.status === 'interrupted';
 
   const handleSubmit = () => {
-    if (!query) return;
+    if (!query || !selectedConnection) return;
     setChatResponse(null);
     const newThreadId = uuidv4();
     setThreadId(newThreadId);
@@ -1019,7 +1040,7 @@ export function AgentTestingPage() {
         query,
         thread_id: newThreadId,
         hitl_enabled: hitlEnabled,
-        allowed_statuses: allowedStatuses.length > 0 ? allowedStatuses : undefined,
+        connection_id: selectedConnection,
       });
     }, 300);
   };
@@ -1030,6 +1051,7 @@ export function AgentTestingPage() {
       chatMutation.mutate({
         thread_id: threadId,
         resume_value: resumeValue !== undefined ? resumeValue : { approved: true },
+        connection_id: selectedConnection,
       });
     }, 300);
   };
@@ -1041,6 +1063,7 @@ export function AgentTestingPage() {
         thread_id: threadId,
         resume_value: { approved: false, feedback, rejection_category: category },
         hitl_enabled: hitlEnabled,
+        connection_id: selectedConnection,
       });
     }, 300);
   };
@@ -1096,8 +1119,10 @@ export function AgentTestingPage() {
       <AgentTestingHeader
         hitlEnabled={hitlEnabled}
         setHitlEnabled={setHitlEnabled}
-        allowedStatuses={allowedStatuses}
-        setAllowedStatuses={setAllowedStatuses}
+        connections={connections}
+        selectedConnection={selectedConnection}
+        setSelectedConnection={setSelectedConnection}
+        loadingConnections={loadingConnections}
       />
 
       <AgentChatInput
@@ -1106,6 +1131,7 @@ export function AgentTestingPage() {
         onSubmit={handleSubmit}
         disabled={isInputDisabled}
         loading={isPendingInitial}
+        submitDisabled={!selectedConnection || !query.trim()}
       />
 
       {chatMutation.isError && (