Update data sources to include descriptions and send them to copilot

This commit is contained in:
akhisud3195 2025-05-09 19:01:06 +05:30
parent 820150641c
commit 6875459327
13 changed files with 218 additions and 17 deletions

View file

@ -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:

View file

@ -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}
"""

View file

@ -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,

View file

@ -2,10 +2,12 @@
import {
convertToCopilotWorkflow, convertToCopilotMessage, convertToCopilotApiMessage,
convertToCopilotApiChatContext, CopilotAPIResponse, CopilotAPIRequest,
CopilotChatContext, CopilotMessage, CopilotAssistantMessage, CopilotWorkflow
CopilotChatContext, CopilotMessage, CopilotAssistantMessage, CopilotWorkflow,
CopilotDataSource
} from "../lib/types/copilot_types";
import {
Workflow} from "../lib/types/workflow_types";
import { DataSource } from "../lib/types/datasource_types";
import { z } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';
import { assert } from "node:console";
@ -18,7 +20,8 @@ export async function getCopilotResponse(
projectId: string,
messages: z.infer<typeof CopilotMessage>[],
current_workflow_config: z.infer<typeof Workflow>,
context: z.infer<typeof CopilotChatContext> | null
context: z.infer<typeof CopilotChatContext> | null,
dataSources?: z.infer<typeof DataSource>[]
): Promise<{
message: z.infer<typeof CopilotAssistantMessage>;
rawRequest: unknown;
@ -35,6 +38,7 @@ 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,
};
console.log(`sending copilot request`, JSON.stringify(request));
@ -96,7 +100,8 @@ export async function getCopilotResponseStream(
projectId: string,
messages: z.infer<typeof CopilotMessage>[],
current_workflow_config: z.infer<typeof Workflow>,
context: z.infer<typeof CopilotChatContext> | null
context: z.infer<typeof CopilotChatContext> | null,
dataSources?: z.infer<typeof DataSource>[]
): Promise<{
streamId: string;
}> {
@ -111,6 +116,7 @@ export async function getCopilotResponseStream(
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,
};
// serialize the request

View file

@ -43,11 +43,13 @@ export async function listDataSources(projectId: string): Promise<WithStringId<z
export async function createDataSource({
projectId,
name,
description,
data,
status = 'pending',
}: {
projectId: string,
name: string,
description?: string,
data: z.infer<typeof DataSource>['data'],
status?: 'pending' | 'ready',
}): Promise<WithStringId<z.infer<typeof DataSource>>> {
@ -57,6 +59,7 @@ export async function createDataSource({
projectId: projectId,
active: true,
name: name,
description,
createdAt: (new Date()).toISOString(),
attempts: 0,
status: status,
@ -354,3 +357,28 @@ export async function getUploadUrlsForFilesDataSource(
return urls;
}
export async function updateDataSource({
projectId,
sourceId,
description,
}: {
projectId: string,
sourceId: string,
description: string,
}) {
await projectAuthCheck(projectId);
await getDataSource(projectId, sourceId);
await dataSourcesCollection.updateOne({
_id: new ObjectId(sourceId),
}, {
$set: {
description,
lastUpdatedAt: (new Date()).toISOString(),
},
$inc: {
version: 1,
},
});
}

View file

@ -3,11 +3,23 @@ import { Workflow } from "./workflow_types";
import { apiV1 } from "rowboat-shared"
import { AgenticAPIChatMessage } from "./agents_api_types";
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 CopilotWorkflow = Workflow.omit({
lastUpdatedAt: true,
projectId: true,
});export const CopilotUserMessage = z.object({
});
export const CopilotUserMessage = z.object({
role: z.literal('user'),
content: z.string(),
});
@ -77,6 +89,7 @@ export const CopilotAPIRequest = z.object({
workflow_schema: z.string(),
current_workflow_config: z.string(),
context: CopilotApiChatContext.nullable(),
dataSources: z.array(CopilotDataSource).optional(),
});
export const CopilotAPIResponse = z.union([
z.object({

View file

@ -2,6 +2,7 @@ import { z } from "zod";
export const DataSource = z.object({
name: z.string(),
description: z.string().optional(),
projectId: z.string(),
active: z.boolean().default(true),
status: z.union([

View file

@ -5,6 +5,7 @@ import { useRef, useState, createContext, useContext, useCallback, forwardRef, u
import { CopilotChatContext } from "../../../lib/types/copilot_types";
import { CopilotMessage } from "../../../lib/types/copilot_types";
import { Workflow } from "@/app/lib/types/workflow_types";
import { DataSource } from "@/app/lib/types/datasource_types";
import { z } from "zod";
import { Action as WorkflowDispatch } from "../workflow/workflow_editor";
import { Panel } from "@/components/common/panel-common";
@ -30,6 +31,7 @@ interface AppProps {
onCopyJson?: (data: { messages: any[] }) => void;
onMessagesChange?: (messages: z.infer<typeof CopilotMessage>[]) => void;
isInitialState?: boolean;
dataSources?: z.infer<typeof DataSource>[];
}
const App = forwardRef<{ handleCopyChat: () => void }, AppProps>(function App({
@ -40,6 +42,7 @@ const App = forwardRef<{ handleCopyChat: () => void }, AppProps>(function App({
onCopyJson,
onMessagesChange,
isInitialState = false,
dataSources,
}, ref) {
const [messages, setMessages] = useState<z.infer<typeof CopilotMessage>[]>([]);
const [discardContext, setDiscardContext] = useState(false);
@ -63,7 +66,8 @@ const App = forwardRef<{ handleCopyChat: () => void }, AppProps>(function App({
} = useCopilot({
projectId,
workflow: workflowRef.current,
context: effectiveContext
context: effectiveContext,
dataSources: dataSources
});
// Store latest start/cancel functions in refs
@ -196,12 +200,14 @@ export function Copilot({
chatContext = undefined,
dispatch,
isInitialState = false,
dataSources,
}: {
projectId: string;
workflow: z.infer<typeof Workflow>;
chatContext?: z.infer<typeof CopilotChatContext>;
dispatch: (action: WorkflowDispatch) => void;
isInitialState?: boolean;
dataSources?: z.infer<typeof DataSource>[];
}) {
const [copilotKey, setCopilotKey] = useState(0);
const [showCopySuccess, setShowCopySuccess] = useState(false);
@ -277,6 +283,7 @@ export function Copilot({
onCopyJson={handleCopyJson}
onMessagesChange={setMessages}
isInitialState={isInitialState}
dataSources={dataSources}
/>
</div>
</Panel>

View file

@ -2,12 +2,14 @@ import { useCallback, useRef, useState } from "react";
import { getCopilotResponseStream } from "@/app/actions/copilot_actions";
import { CopilotMessage } from "@/app/lib/types/copilot_types";
import { Workflow } from "@/app/lib/types/workflow_types";
import { DataSource } from "@/app/lib/types/datasource_types";
import { z } from "zod";
interface UseCopilotParams {
projectId: string;
workflow: z.infer<typeof Workflow>;
context: any;
dataSources?: z.infer<typeof DataSource>[];
}
interface UseCopilotResult {
@ -21,7 +23,7 @@ interface UseCopilotResult {
cancel: () => void;
}
export function useCopilot({ projectId, workflow, context }: UseCopilotParams): UseCopilotResult {
export function useCopilot({ projectId, workflow, context, dataSources }: UseCopilotParams): UseCopilotResult {
const [streamingResponse, setStreamingResponse] = useState('');
const [loading, setLoading] = useState(false);
const [error, setError] = useState<string | null>(null);
@ -41,7 +43,7 @@ export function useCopilot({ projectId, workflow, context }: UseCopilotParams):
setLoading(true);
try {
const res = await getCopilotResponseStream(projectId, messages, workflow, context || null);
const res = await getCopilotResponseStream(projectId, messages, workflow, context || null, dataSources);
const eventSource = new EventSource(`/api/copilot-stream-response/${res.streamId}`);
eventSource.onmessage = (event) => {
@ -71,7 +73,7 @@ export function useCopilot({ projectId, workflow, context }: UseCopilotParams):
setError('Failed to initiate stream');
setLoading(false);
}
}, [projectId, workflow, context]);
}, [projectId, workflow, context, dataSources]);
const cancel = useCallback(() => {
cancelRef.current?.();

View file

@ -10,12 +10,14 @@ import { DataSourceIcon } from "../../../../lib/components/datasource-icon";
import { z } from "zod";
import { ScrapeSource } from "../components/scrape-source";
import { FilesSource } from "../components/files-source";
import { getDataSource } from "../../../../actions/datasource_actions";
import { getDataSource, updateDataSource } from "../../../../actions/datasource_actions";
import { TextSource } from "../components/text-source";
import { Panel } from "@/components/common/panel-common";
import { Section, SectionRow, SectionLabel, SectionContent } from "../components/section";
import Link from "next/link";
import { BackIcon } from "../../../../lib/components/icons";
import { Textarea } from "@/components/ui/textarea";
import { CheckIcon } from "lucide-react";
export function SourcePage({
sourceId,
@ -26,6 +28,7 @@ export function SourcePage({
}) {
const [source, setSource] = useState<WithStringId<z.infer<typeof DataSource>> | null>(null);
const [isLoading, setIsLoading] = useState(true);
const [showSaveSuccess, setShowSaveSuccess] = useState(false);
async function handleReload() {
setIsLoading(true);
@ -122,6 +125,57 @@ export function SourcePage({
</SectionContent>
</SectionRow>
<SectionRow>
<SectionLabel>Name</SectionLabel>
<SectionContent>
<div className="text-sm text-gray-900 dark:text-gray-100">
{source.name}
</div>
</SectionContent>
</SectionRow>
<SectionRow>
<SectionLabel className="pt-3">Description</SectionLabel>
<SectionContent>
<form
action={async (formData: FormData) => {
const description = formData.get('description') as string;
await updateDataSource({
projectId,
sourceId,
description,
});
handleReload();
setShowSaveSuccess(true);
setTimeout(() => setShowSaveSuccess(false), 2000);
}}
className="w-full"
>
<Textarea
name="description"
defaultValue={source.description || ''}
placeholder="Add a description for this data source"
rows={2}
className="w-full rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
<div className="flex items-center gap-2 mt-2">
<button
type="submit"
className="text-sm text-blue-600 hover:text-blue-700 dark:text-blue-400 dark:hover:text-blue-300"
>
Save
</button>
{showSaveSuccess && (
<div className="flex items-center gap-1 text-sm text-green-600 dark:text-green-400">
<CheckIcon className="w-4 h-4" />
<span>Saved</span>
</div>
)}
</div>
</form>
</SectionContent>
</SectionRow>
<SectionRow>
<SectionLabel>Type</SectionLabel>
<SectionContent>

View file

@ -29,15 +29,15 @@ export function Section({ title, description, children, className }: SectionProp
export function SectionRow({ children, className }: { children: ReactNode; className?: string }) {
return (
<div className={`flex items-center gap-6 ${className || ''}`}>
<div className={`flex items-start gap-6 py-1 ${className || ''}`}>
{children}
</div>
);
}
export function SectionLabel({ children }: { children: ReactNode }) {
export function SectionLabel({ children, className }: { children: ReactNode; className?: string }) {
return (
<div className="w-24 flex-shrink-0 text-sm text-gray-500 dark:text-gray-400">
<div className={`w-24 flex-shrink-0 text-sm text-gray-500 dark:text-gray-400 ${className || ''}`}>
{children}
</div>
);

View file

@ -57,6 +57,7 @@ export function Form({
const source = await createDataSource({
projectId,
name: formData.get('name') as string,
description: formData.get('description') as string,
data: {
type: 'urls',
},
@ -85,6 +86,7 @@ export function Form({
const source = await createDataSource({
projectId,
name: formData.get('name') as string,
description: formData.get('description') as string,
data: {
type: formData.get('type') as 'files_local' | 'files_s3',
},
@ -98,6 +100,7 @@ export function Form({
const source = await createDataSource({
projectId,
name: formData.get('name') as string,
description: formData.get('description') as string,
data: {
type: 'text',
},
@ -167,6 +170,17 @@ export function Form({
className="rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
</div>
<div className="space-y-2">
<label className="text-xs font-medium uppercase tracking-wider text-gray-500 dark:text-gray-400">
Description
</label>
<Textarea
name="description"
placeholder="e.g. A collection of help articles from our documentation"
rows={2}
className="rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
</div>
<div className="bg-gray-50 dark:bg-gray-800/50 rounded-lg p-4 border border-gray-200 dark:border-gray-700">
<div className="flex items-center gap-2 mb-2 text-gray-700 dark:text-gray-300">
<svg
@ -222,6 +236,17 @@ export function Form({
className="rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
</div>
<div className="space-y-2">
<label className="text-xs font-medium uppercase tracking-wider text-gray-500 dark:text-gray-400">
Description
</label>
<Textarea
name="description"
placeholder="e.g. A collection of documentation files"
rows={2}
className="rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
</div>
<div className="bg-gray-50 dark:bg-gray-800/50 rounded-lg p-4 border border-gray-200 dark:border-gray-700">
<div className="flex items-center gap-2 mb-2 text-gray-700 dark:text-gray-300">
<svg
@ -281,6 +306,17 @@ export function Form({
className="rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
</div>
<div className="space-y-2">
<label className="text-xs font-medium uppercase tracking-wider text-gray-500 dark:text-gray-400">
Description
</label>
<Textarea
name="description"
placeholder="e.g. A collection of documentation for our product"
rows={2}
className="rounded-lg p-3 border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-800 hover:bg-gray-50 dark:hover:bg-gray-750 focus:shadow-inner focus:ring-2 focus:ring-indigo-500/20 dark:focus:ring-indigo-400/20 placeholder:text-gray-400 dark:placeholder:text-gray-500"
/>
</div>
<FormStatusButton
props={{
type: "submit",

View file

@ -1033,6 +1033,7 @@ export function WorkflowEditor({
} : undefined
}
isInitialState={isInitialState}
dataSources={dataSources}
/>
</ResizablePanel>
</>