Add data source ids to copilot request

apps/copilot/app.py

@@ -1,5 +1,5 @@
 from flask import Flask, request, jsonify, Response, stream_with_context
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, ValidationError, Field
 from typing import List, Optional
 from copilot import UserMessage, AssistantMessage, get_response
 from streaming import get_streaming_response
@@ -10,6 +10,7 @@ from copilot import copilot_instructions_edit_agent
 import json
 
 class DataSource(BaseModel):
+    id: str = Field(alias='_id')
     name: str
     description: Optional[str] = None
     active: bool = True
@@ -17,6 +18,9 @@ class DataSource(BaseModel):
     error: Optional[str] = None
     data: dict  # The discriminated union based on type
 
+    class Config:
+        populate_by_name = True
+
 class ApiRequest(BaseModel):
     messages: List[UserMessage | AssistantMessage]
     workflow_schema: str
@@ -61,7 +65,10 @@ def health():
 @require_api_key
 def chat_stream():
     try:
-        request_data = ApiRequest(**request.json)
+        raw_data = request.json
+        print(f"Raw request JSON: {json.dumps(raw_data)}")
+        
+        request_data = ApiRequest(**raw_data)
         print(f"received /chat_stream request: {request_data}")
         validate_request(request_data)
 

apps/copilot/copilot.py

@@ -16,6 +16,7 @@ class AssistantMessage(BaseModel):
     content: str
 
 class DataSource(BaseModel):
+    _id: str
     name: str
     description: Optional[str] = None
     active: bool = True

apps/copilot/streaming.py

@@ -1,6 +1,6 @@
 from openai import OpenAI
 from flask import Flask, request, jsonify, Response, stream_with_context
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, ValidationError, Field
 from typing import List, Dict, Any, Literal, Optional
 import json
 from lib import AgentContext, PromptContext, ToolContext, ChatContext
@@ -16,6 +16,7 @@ class AssistantMessage(BaseModel):
     content: str
 
 class DataSource(BaseModel):
+    id: str = Field(alias='_id')
     name: str
     description: Optional[str] = None
     active: bool = True
@@ -23,6 +24,9 @@ class DataSource(BaseModel):
     error: Optional[str] = None
     data: dict  # The discriminated union based on type
 
+    class Config:
+        populate_by_name = True
+
 with open('copilot_multi_agent.md', 'r', encoding='utf-8') as file:
     copilot_instructions_multi_agent = file.read()
 
@@ -81,6 +85,7 @@ def get_streaming_response(
     data_sources_prompt = ""
     if dataSources:
         print(f"Data sources found at project level: {dataSources}")
+        print(f"Data source IDs: {[ds.id for ds in dataSources]}")
         data_sources_prompt = f"""
 **NOTE**: The following data sources are available:
 ```json
@@ -136,6 +141,8 @@ def create_app():
             if not request_data or 'messages' not in request_data:
                 return jsonify({'error': 'No messages provided'}), 400
 
+            print(f"Raw request data: {request_data}")
+
             messages = [
                 UserMessage(**msg) if msg['role'] == 'user' else AssistantMessage(**msg)
                 for msg in request_data['messages']
@@ -144,13 +151,19 @@ def create_app():
             workflow_schema = request_data.get('workflow_schema', '')
             current_workflow_config = request_data.get('current_workflow_config', '')
             context = None  # You can add context handling if needed
+            dataSources = None
+            if 'dataSources' in request_data and request_data['dataSources']:
+                print(f"Raw dataSources from request: {request_data['dataSources']}")
+                dataSources = [DataSource(**ds) for ds in request_data['dataSources']]
+                print(f"Parsed dataSources: {dataSources}")
 
             def generate():
                 stream = get_streaming_response(
                     messages=messages,
                     workflow_schema=workflow_schema,
                     current_workflow_config=current_workflow_config,
-                    context=context
+                    context=context,
+                    dataSources=dataSources
                 )
 
                 for chunk in stream:

apps/rowboat/app/actions/copilot_actions.ts

@@ -38,7 +38,24 @@ export async function getCopilotResponse(
         workflow_schema: JSON.stringify(zodToJsonSchema(CopilotWorkflow)),
         current_workflow_config: JSON.stringify(convertToCopilotWorkflow(current_workflow_config)),
         context: context ? convertToCopilotApiChatContext(context) : null,
-        dataSources: dataSources ? dataSources.map(ds => CopilotDataSource.parse(ds)) : undefined,
+        dataSources: dataSources ? dataSources.map(ds => {
+            console.log('Original data source:', JSON.stringify(ds));
+            // First parse to validate, then ensure _id is included
+            CopilotDataSource.parse(ds); // validate but don't use the result
+            // Cast to any to handle the WithStringId type
+            const withId = ds as any;
+            const result = {
+                _id: withId._id,
+                name: withId.name,
+                description: withId.description,
+                active: withId.active,
+                status: withId.status,
+                error: withId.error,
+                data: withId.data
+            };
+            console.log('Processed data source:', JSON.stringify(result));
+            return result;
+        }) : undefined,
     };
     console.log(`sending copilot request`, JSON.stringify(request));
 

apps/rowboat/app/lib/types/copilot_types.ts

@@ -6,14 +6,33 @@ import { convertToAgenticAPIChatMessages } from "./agents_api_types";
 import { DataSource } from "./datasource_types";
 
 // Create a filtered version of DataSource for copilot
-export const CopilotDataSource = DataSource.omit({
-    projectId: true,
-    version: true,
-    attempts: true,
-    createdAt: true,
-    lastUpdatedAt: true,
-    pendingRefresh: true,
-});
+export const CopilotDataSource = z.object({
+    _id: z.string(),
+    name: z.string(),
+    description: z.string().optional(),
+    active: z.boolean().default(true),
+    status: z.union([
+        z.literal('pending'),
+        z.literal('ready'),
+        z.literal('error'),
+        z.literal('deleted'),
+    ]),
+    error: z.string().optional(),
+    data: z.discriminatedUnion('type', [
+        z.object({
+            type: z.literal('urls'),
+        }),
+        z.object({
+            type: z.literal('files_local'),
+        }),
+        z.object({
+            type: z.literal('files_s3'),
+        }),
+        z.object({
+            type: z.literal('text'),
+        })
+    ]),
+}).passthrough();
 
 export const CopilotWorkflow = Workflow.omit({
     lastUpdatedAt: true,