backend/assistants_web: Move Default Agent usage to backend (#842)

* Add router for default agent

* Working app with default agent

* Lint

* Working app and tests

* add back global exc handler

* prettify coral

* PR review

Author: tianjing-li

src/backend/config/default_agent.py

@@ -0,0 +1,33 @@
+import datetime
+
+from backend.config.deployments import ModelDeploymentName
+from backend.config.tools import Tool
+from backend.schemas.agent import AgentPublic
+
+DEFAULT_AGENT_ID = "default"
+DEFAULT_DEPLOYMENT = ModelDeploymentName.CoherePlatform
+DEFAULT_MODEL = "command-r-plus"
+
+def get_default_agent() -> AgentPublic:
+    return AgentPublic(
+        id=DEFAULT_AGENT_ID,
+        name='Command R+',
+        description='Ask questions and get answers based on your files.',
+        created_at=datetime.datetime.now(),
+        updated_at=datetime.datetime.now(),
+        preamble="",
+        version=1,
+        temperature=0.3,
+        tools=[
+            Tool.Read_File.value.ID,
+            Tool.Search_File.value.ID,
+            Tool.Python_Interpreter.value.ID,
+            Tool.Hybrid_Web_Search.value.ID,
+        ],
+        tools_metadata=[],
+        deployment=DEFAULT_DEPLOYMENT,
+        model=DEFAULT_MODEL,
+        user_id='',
+        organization_id=None,
+        is_private=False,
+    )

src/backend/main.py

@@ -91,7 +91,6 @@ def create_app():
 
 app = create_app()
 
-
 @app.exception_handler(Exception)
 async def validation_exception_handler(request: Request, exc: Exception):
     ctx = get_context(request)
@@ -115,7 +114,6 @@ async def validation_exception_handler(request: Request, exc: Exception):
         },
     )
 
-
 @app.on_event("startup")
 async def startup_event():
     """

src/backend/routers/agent.py

@@ -5,6 +5,7 @@
 from fastapi import File as RequestFile
 from fastapi import UploadFile as FastAPIUploadFile
 
+from backend.config.default_agent import DEFAULT_AGENT_ID, get_default_agent
 from backend.config.routers import RouterName
 from backend.crud import agent as agent_crud
 from backend.crud import agent_tool_metadata as agent_tool_metadata_crud
@@ -71,9 +72,9 @@
     ],
 )
 async def create_agent(
-        session: DBSessionDep,
-        agent: CreateAgentRequest,
-        ctx: Context = Depends(get_context),
+    session: DBSessionDep,
+    agent: CreateAgentRequest,
+    ctx: Context = Depends(get_context),
 ) -> AgentPublic:
     """
     Create an agent.
@@ -127,13 +128,13 @@ async def create_agent(
 
 @router.get("", response_model=list[AgentPublic])
 async def list_agents(
-        *,
-        offset: int = 0,
-        limit: int = 100,
-        session: DBSessionDep,
-        visibility: AgentVisibility = AgentVisibility.ALL,
-        organization_id: Optional[str] = None,
-        ctx: Context = Depends(get_context),
+    *,
+    offset: int = 0,
+    limit: int = 100,
+    session: DBSessionDep,
+    visibility: AgentVisibility = AgentVisibility.ALL,
+    organization_id: Optional[str] = None,
+    ctx: Context = Depends(get_context),
 ) -> list[AgentPublic]:
     """
     List all agents.
@@ -163,6 +164,8 @@ async def list_agents(
             visibility=visibility,
             organization_id=organization_id,
         )
+        # Tradeoff: This appends the default Agent regardless of pagination
+        agents.append(get_default_agent())
         return agents
     except Exception as e:
         logger.exception(event=e)
@@ -171,8 +174,8 @@ async def list_agents(
 
 @router.get("/{agent_id}", response_model=AgentPublic)
 async def get_agent_by_id(
-        agent_id: str, session: DBSessionDep, ctx: Context = Depends(get_context)
-) -> Agent:
+    agent_id: str, session: DBSessionDep, ctx: Context = Depends(get_context)
+) -> AgentPublic:
     """
     Args:
         agent_id (str): Agent ID.
@@ -189,7 +192,11 @@ async def get_agent_by_id(
     agent = None
 
     try:
-        agent = agent_crud.get_agent_by_id(session, agent_id, user_id)
+        # Intentionally not adding Default Agent to DB so it's more flexible
+        if agent_id == DEFAULT_AGENT_ID:
+            agent = get_default_agent()
+        else:
+            agent = agent_crud.get_agent_by_id(session, agent_id, user_id)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 

src/backend/tests/unit/routers/test_agent.py

@@ -4,6 +4,7 @@
 from fastapi.testclient import TestClient
 from sqlalchemy.orm import Session
 
+from backend.config.default_agent import DEFAULT_AGENT_ID
 from backend.config.deployments import ModelDeploymentName
 from backend.config.tools import Tool
 from backend.crud import deployment as deployment_crud
@@ -17,6 +18,9 @@
     and os.environ.get("COHERE_API_KEY") != ""
 )
 
+def filter_default_agent(agents: list) -> list:
+    return [agent for agent in agents if agent.get("id") != DEFAULT_AGENT_ID]
+
 def test_create_agent_missing_name(
     session_client: TestClient, session: Session, user
 ) -> None:
@@ -159,35 +163,34 @@ def test_create_existing_agent(
     assert response.json() == {"detail": "Agent test agent already exists."}
 
 
-def test_list_agents_empty(session_client: TestClient, session: Session) -> None:
-    # Delete default agent
-    session.query(Agent).delete()
+def test_list_agents_empty_returns_default_agent(session_client: TestClient, session: Session) -> None:
     response = session_client.get("/v1/agents", headers={"User-Id": "123"})
     assert response.status_code == 200
     response_agents = response.json()
-    assert len(response_agents) == 0
+    # Returns default agent
+    assert len(response_agents) == 1
 
 
 def test_list_agents(session_client: TestClient, session: Session, user) -> None:
-    session.query(Agent).delete()
-    for _ in range(3):
+    num_agents = 3
+    for _ in range(num_agents):
         _ = get_factory("Agent", session).create(user=user)
 
     response = session_client.get("/v1/agents", headers={"User-Id": user.id})
     assert response.status_code == 200
-    response_agents = response.json()
-    assert len(response_agents) == 3
+    response_agents = filter_default_agent(response.json())
+    assert len(response_agents) == num_agents
 
 
 def test_list_organization_agents(
     session_client: TestClient,
     session: Session,
     user,
 ) -> None:
-    session.query(Agent).delete()
+    num_agents = 3
     organization = get_factory("Organization", session).create()
     organization1 = get_factory("Organization", session).create()
-    for i in range(3):
+    for i in range(num_agents):
         _ = get_factory("Agent", session).create(
             user=user,
             organization_id=organization.id,
@@ -201,9 +204,9 @@ def test_list_organization_agents(
         "/v1/agents", headers={"User-Id": user.id, "Organization-Id": organization.id}
     )
     assert response.status_code == 200
-    response_agents = response.json()
+    response_agents = filter_default_agent(response.json())
     agents = sorted(response_agents, key=lambda x: x["name"])
-    for i in range(3):
+    for i in range(num_agents):
         assert agents[i]["name"] == f"agent-{i}-{organization.id}"
 
 
@@ -212,10 +215,10 @@ def test_list_organization_agents_query_param(
     session: Session,
     user,
 ) -> None:
-    session.query(Agent).delete()
+    num_agents = 3
     organization = get_factory("Organization", session).create()
     organization1 = get_factory("Organization", session).create()
-    for i in range(3):
+    for i in range(num_agents):
         _ = get_factory("Agent", session).create(
             user=user, organization_id=organization.id
         )
@@ -230,9 +233,9 @@ def test_list_organization_agents_query_param(
         headers={"User-Id": user.id, "Organization-Id": organization.id},
     )
     assert response.status_code == 200
-    response_agents = response.json()
+    response_agents = filter_default_agent(response.json())
     agents = sorted(response_agents, key=lambda x: x["name"])
-    for i in range(3):
+    for i in range(num_agents):
         assert agents[i]["name"] == f"agent-{i}-{organization1.id}"
 
 
@@ -263,7 +266,7 @@ def test_list_private_agents(
     )
 
     assert response.status_code == 200
-    response_agents = response.json()
+    response_agents = filter_default_agent(response.json())
 
     # Only the agents created by user should be returned
     assert len(response_agents) == 3
@@ -282,7 +285,7 @@ def test_list_public_agents(session_client: TestClient, session: Session, user)
     )
 
     assert response.status_code == 200
-    response_agents = response.json()
+    response_agents = filter_default_agent(response.json())
 
     # Only the agents created by user should be returned
     assert len(response_agents) == 2
@@ -319,14 +322,14 @@ def test_list_agents_with_pagination(
         "/v1/agents?limit=3&offset=2", headers={"User-Id": user.id}
     )
     assert response.status_code == 200
-    response_agents = response.json()
+    response_agents = filter_default_agent(response.json())
     assert len(response_agents) == 3
 
     response = session_client.get(
         "/v1/agents?limit=2&offset=4", headers={"User-Id": user.id}
     )
     assert response.status_code == 200
-    response_agents = response.json()
+    response_agents = filter_default_agent(response.json())
     assert len(response_agents) == 1
 
 

src/interfaces/assistants_web/src/app/(main)/(chat)/c/[conversationId]/page.tsx

@@ -2,7 +2,7 @@ import { HydrationBoundary, QueryClient, dehydrate } from '@tanstack/react-query
 import { NextPage } from 'next';
 
 import Chat from '@/app/(main)/(chat)/Chat';
-import { BASE_AGENT } from '@/constants';
+import { DEFAULT_AGENT_ID } from '@/constants';
 import { getCohereServerClient } from '@/server/cohereServerClient';
 
 type Props = {
@@ -23,8 +23,8 @@ const Page: NextPage<Props> = async ({ params }) => {
         cohereServerClient.getConversation({ conversationId: params.conversationId }),
     }),
     queryClient.prefetchQuery({
-      queryKey: ['agent', null],
-      queryFn: () => BASE_AGENT,
+      queryKey: ['agent', DEFAULT_AGENT_ID],
+      queryFn: () => cohereServerClient.getDefaultAgent(),
     }),
   ]);
 

src/interfaces/assistants_web/src/app/(main)/(chat)/page.tsx

@@ -2,14 +2,16 @@ import { HydrationBoundary, QueryClient, dehydrate } from '@tanstack/react-query
 import { NextPage } from 'next';
 
 import Chat from '@/app/(main)/(chat)/Chat';
-import { BASE_AGENT } from '@/constants';
+import { DEFAULT_AGENT_ID } from '@/constants';
+import { getCohereServerClient } from '@/server/cohereServerClient';
 
 const Page: NextPage = async () => {
   const queryClient = new QueryClient();
+  const cohereServerClient = getCohereServerClient();
 
   await queryClient.prefetchQuery({
-    queryKey: ['agent', null],
-    queryFn: () => BASE_AGENT,
+    queryKey: ['agent', DEFAULT_AGENT_ID],
+    queryFn: () => cohereServerClient.getDefaultAgent(),
   });
 
   return (

src/interfaces/assistants_web/src/app/(main)/discover/DiscoverAgentCard.tsx

@@ -7,7 +7,7 @@ import { DeleteAgent } from '@/components/Modals/DeleteAgent';
 import { CoralLogo, KebabMenu, Text } from '@/components/UI';
 import { useContextStore } from '@/context';
 import { useBrandedColors, useSession } from '@/hooks';
-import { checkIsBaseAgent, cn } from '@/utils';
+import { checkIsDefaultAgent, cn } from '@/utils';
 
 type Props = {
   agent?: AgentPublic;
@@ -17,11 +17,11 @@ type Props = {
  * @description renders a card for an agent with the agent's name, description
  */
 export const DiscoverAgentCard: React.FC<Props> = ({ agent }) => {
-  const isBaseAgent = checkIsBaseAgent(agent);
+  const isDefaultAgent = checkIsDefaultAgent(agent);
   const { bg, contrastText, contrastFill } = useBrandedColors(agent?.id);
   const session = useSession();
   const isCreator = agent?.user_id === session.userId;
-  const createdBy = isBaseAgent ? 'COHERE' : isCreator ? 'YOU' : 'TEAM';
+  const createdBy = isDefaultAgent ? 'COHERE' : isCreator ? 'YOU' : 'TEAM';
 
   const { open, close } = useContextStore();
 
@@ -36,7 +36,7 @@ export const DiscoverAgentCard: React.FC<Props> = ({ agent }) => {
   return (
     <Link
       className="flex overflow-x-hidden rounded-lg border border-volcanic-800 bg-volcanic-950 p-4 transition-colors duration-300 hover:bg-marble-950 dark:border-volcanic-300 dark:bg-volcanic-150 dark:hover:bg-volcanic-100"
-      href={isBaseAgent ? '/' : `/a/${agent?.id}`}
+      href={isDefaultAgent ? '/' : `/a/${agent?.id}`}
     >
       <div className="flex h-full flex-grow flex-col items-start gap-y-2 overflow-x-hidden">
         <div className="flex w-full items-center gap-x-2">
@@ -46,7 +46,7 @@ export const DiscoverAgentCard: React.FC<Props> = ({ agent }) => {
               bg
             )}
           >
-            {isBaseAgent ? (
+            {isDefaultAgent ? (
               <CoralLogo className={contrastFill} />
             ) : (
               <Text className={cn('uppercase', contrastText)} styleAs="p-lg">

src/interfaces/assistants_web/src/app/(main)/layout.tsx

@@ -6,7 +6,7 @@ import { redirect } from 'next/navigation';
 import { Swipeable } from '@/components/Global';
 import { HotKeys } from '@/components/HotKeys';
 import { SideNavPanel } from '@/components/SideNavPanel';
-import { COOKIE_KEYS, DEFAULT_AGENT_TOOLS } from '@/constants';
+import { BACKGROUND_TOOLS, COOKIE_KEYS } from '@/constants';
 import { getCohereServerClient } from '@/server/cohereServerClient';
 
 const MainLayout: NextPage<React.PropsWithChildren> = async ({ children }) => {
@@ -35,7 +35,7 @@ const MainLayout: NextPage<React.PropsWithChildren> = async ({ children }) => {
       queryKey: ['tools'],
       queryFn: async () => {
         const tools = await cohereServerClient.listTools({});
-        return tools.filter((tool) => !DEFAULT_AGENT_TOOLS.includes(tool.name ?? ''));
+        return tools.filter((tool) => !BACKGROUND_TOOLS.includes(tool.name ?? ''));
       },
     }),
     queryClient.prefetchQuery({

src/interfaces/assistants_web/src/app/(main)/new/CreateAgent.tsx

@@ -9,8 +9,8 @@ import { AgentSettingsFields, AgentSettingsForm } from '@/components/AgentSettin
 import { MobileHeader } from '@/components/Global';
 import { Button, Icon, Text } from '@/components/UI';
 import {
+  BACKGROUND_TOOLS,
   DEFAULT_AGENT_MODEL,
-  DEFAULT_AGENT_TOOLS,
   DEFAULT_PREAMBLE,
   DEPLOYMENT_COHERE_PLATFORM,
 } from '@/constants';
@@ -23,7 +23,7 @@ const DEFAULT_FIELD_VALUES = {
   preamble: DEFAULT_PREAMBLE,
   deployment: DEPLOYMENT_COHERE_PLATFORM,
   model: DEFAULT_AGENT_MODEL,
-  tools: DEFAULT_AGENT_TOOLS,
+  tools: BACKGROUND_TOOLS,
   is_private: false,
 };
 /**

src/interfaces/assistants_web/src/cohere-client/client.ts

@@ -17,6 +17,7 @@ import {
   UpdateConversationRequest,
   UpdateDeploymentEnv,
 } from '@/cohere-client';
+import { DEFAULT_AGENT_ID } from '@/constants';
 
 import { mapToChatRequest } from './mappings';
 
@@ -285,6 +286,10 @@ export class CohereClient {
     // this.cohereService.default.oidcAuthorizeV1OidcAuthGet();
   }
 
+  public getDefaultAgent() {
+    return this.cohereService.default.getAgentByIdV1AgentsAgentIdGet({ agentId: DEFAULT_AGENT_ID });
+  }
+
   public getAgent(agentId: string) {
     return this.cohereService.default.getAgentByIdV1AgentsAgentIdGet({ agentId });
   }