Update data sources to include descriptions and send them to copilot

apps/copilot/app.py

@@ -1,6 +1,6 @@
 from flask import Flask, request, jsonify, Response, stream_with_context
 from pydantic import BaseModel, ValidationError
-from typing import List
+from typing import List, Optional
 from copilot import UserMessage, AssistantMessage, get_response
 from streaming import get_streaming_response
 from lib import AgentContext, PromptContext, ToolContext, ChatContext
@@ -9,11 +9,20 @@ from functools import wraps
 from copilot import copilot_instructions_edit_agent
 import json
 
+class DataSource(BaseModel):
+    name: str
+    description: Optional[str] = None
+    active: bool = True
+    status: str  # 'pending' | 'ready' | 'error' | 'deleted'
+    error: Optional[str] = None
+    data: dict  # The discriminated union based on type
+
 class ApiRequest(BaseModel):
     messages: List[UserMessage | AssistantMessage]
     workflow_schema: str
     current_workflow_config: str
     context: AgentContext | PromptContext | ToolContext | ChatContext | None = None
+    dataSources: Optional[List[DataSource]] = None
 
 class ApiResponse(BaseModel):
     response: str
@@ -61,7 +70,8 @@ def chat_stream():
                 messages=request_data.messages,
                 workflow_schema=request_data.workflow_schema,
                 current_workflow_config=request_data.current_workflow_config,
-                context=request_data.context
+                context=request_data.context,
+                dataSources=request_data.dataSources
             )
 
             for chunk in stream:

apps/copilot/copilot.py

@@ -1,7 +1,7 @@
 from openai import OpenAI
 from flask import Flask, request, jsonify
 from pydantic import BaseModel, ValidationError
-from typing import List, Dict, Any, Literal
+from typing import List, Dict, Any, Literal, Optional
 import json
 from lib import AgentContext, PromptContext, ToolContext, ChatContext
 from client import PROVIDER_COPILOT_MODEL
@@ -15,6 +15,14 @@ class AssistantMessage(BaseModel):
     role: Literal["assistant"]
     content: str
 
+class DataSource(BaseModel):
+    name: str
+    description: Optional[str] = None
+    active: bool = True
+    status: str  # 'pending' | 'ready' | 'error' | 'deleted'
+    error: Optional[str] = None
+    data: dict  # The discriminated union based on type
+
 with open('copilot_edit_agent.md', 'r', encoding='utf-8') as file:
     copilot_instructions_edit_agent = file.read()
 
@@ -23,6 +31,7 @@ def get_response(
         workflow_schema: str,
         current_workflow_config: str,
         context: AgentContext | PromptContext | ToolContext | ChatContext | None = None,
+        dataSources: Optional[List[DataSource]] = None,
         copilot_instructions: str = copilot_instructions_edit_agent
     ) -> str:
     # if context is provided, create a prompt for the context
@@ -53,6 +62,16 @@ def get_response(
     else:
         context_prompt = ""
 
+    # Add dataSources to the context if provided
+    data_sources_prompt = ""
+    if dataSources:
+        data_sources_prompt = f"""
+**NOTE**: The following data sources are available:
+```json
+{json.dumps([ds.model_dump() for ds in dataSources])}
+```
+"""
+
     # add the workflow schema to the system prompt
     sys_prompt = copilot_instructions.replace("{workflow_schema}", workflow_schema)
 
@@ -66,6 +85,7 @@ The current workflow config is:
 ```
 
 {context_prompt}
+{data_sources_prompt}
 
 User: {last_message.content}
 """

apps/copilot/streaming.py

@@ -1,7 +1,7 @@
 from openai import OpenAI
 from flask import Flask, request, jsonify, Response, stream_with_context
 from pydantic import BaseModel, ValidationError
-from typing import List, Dict, Any, Literal
+from typing import List, Dict, Any, Literal, Optional
 import json
 from lib import AgentContext, PromptContext, ToolContext, ChatContext
 from client import PROVIDER_COPILOT_MODEL, PROVIDER_DEFAULT_MODEL
@@ -15,6 +15,14 @@ class AssistantMessage(BaseModel):
     role: Literal["assistant"]
     content: str
 
+class DataSource(BaseModel):
+    name: str
+    description: Optional[str] = None
+    active: bool = True
+    status: str  # 'pending' | 'ready' | 'error' | 'deleted'
+    error: Optional[str] = None
+    data: dict  # The discriminated union based on type
+
 with open('copilot_multi_agent.md', 'r', encoding='utf-8') as file:
     copilot_instructions_multi_agent = file.read()
 
@@ -39,6 +47,7 @@ def get_streaming_response(
         workflow_schema: str,
         current_workflow_config: str,
         context: AgentContext | PromptContext | ToolContext | ChatContext | None = None,
+        dataSources: Optional[List[DataSource]] = None,
 ) -> Any:
     # if context is provided, create a prompt for the context
     if context:
@@ -68,6 +77,19 @@ def get_streaming_response(
     else:
         context_prompt = ""
 
+    # Add dataSources to the context if provided
+    data_sources_prompt = ""
+    if dataSources:
+        print(f"Data sources found at project level: {dataSources}")
+        data_sources_prompt = f"""
+**NOTE**: The following data sources are available:
+```json
+{json.dumps([ds.model_dump() for ds in dataSources])}
+```
+"""
+    else:
+        print("No data sources found at project level")
+
     # add the workflow schema to the system prompt
     sys_prompt = streaming_instructions.replace("{workflow_schema}", workflow_schema)
 
@@ -84,6 +106,7 @@ The current workflow config is:
 ```
 
 {context_prompt}
+{data_sources_prompt}
 
 User: {last_message.content}
 """
@@ -91,7 +114,7 @@ User: {last_message.content}
     updated_msgs = [{"role": "system", "content": sys_prompt}] + [
         message.model_dump() for message in messages
     ]
-
+    
     return completions_client.chat.completions.create(
         model=PROVIDER_COPILOT_MODEL,
         messages=updated_msgs,