refactor: add agent-specific routes (#23)

Author: enitrat

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ permissions:
 
 jobs:
   build:
-    name: Build
+    name: Build and Test
     runs-on: ubuntu-latest
 
     steps:
@@ -51,3 +51,9 @@ jobs:
 
       - name: Build
         run: turbo build
+
+      - name: Setup test environment
+        run: cp packages/agents/sample.config.toml packages/agents/config.toml
+
+      - name: Test
+        run: turbo test

packages/agents/__tests__/answerGenerator.test.ts

@@ -13,16 +13,24 @@ import { IterableReadableStream } from '@langchain/core/utils/stream';
 import { BaseMessage, BaseMessageChunk } from '@langchain/core/messages';
 import { BaseLanguageModelInput } from '@langchain/core/language_models/base';
 
-
 // Mock the formatChatHistoryAsString utility
 jest.mock('../src/utils/index', () => ({
   __esModule: true,
-  formatChatHistoryAsString: jest.fn().mockImplementation(() => 'mocked chat history'),
+  formatChatHistoryAsString: jest
+    .fn()
+    .mockImplementation(() => 'mocked chat history'),
   logger: {
     info: jest.fn(),
     debug: jest.fn(),
     error: jest.fn(),
-  }
+  },
+  TokenTracker: {
+    trackFullUsage: jest.fn().mockReturnValue({
+      promptTokens: 100,
+      responseTokens: 50,
+      totalTokens: 150,
+    }),
+  },
 }));
 
 // No need to separately mock the logger since it's now mocked as part of utils/index

packages/agents/__tests__/queryProcessor.test.ts

@@ -1,10 +1,6 @@
 import { QueryProcessor } from '../src/core/pipeline/queryProcessor';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
-import {
-  RagInput,
-  RagSearchConfig,
-  DocumentSource,
-} from '../src/types/index';
+import { RagInput, RagSearchConfig, DocumentSource } from '../src/types/index';
 import { mockDeep, MockProxy } from 'jest-mock-extended';
 import { AIMessage } from '@langchain/core/messages';
 
@@ -15,14 +11,23 @@ jest.mock('../src/utils/index', () => ({
     debug: jest.fn(),
     error: jest.fn(),
   },
-  formatChatHistoryAsString: jest.fn((history) => 
-    history.map((message) => `${message._getType()}: ${message.content}`).join('\n')
+  formatChatHistoryAsString: jest.fn((history) =>
+    history
+      .map((message) => `${message._getType()}: ${message.content}`)
+      .join('\n'),
   ),
   parseXMLContent: jest.fn((xml, tag) => {
     const regex = new RegExp(`<${tag}>(.*?)</${tag}>`, 'gs');
     const matches = [...xml.matchAll(regex)];
     return matches.map((match) => match[1].trim());
   }),
+  TokenTracker: {
+    trackFullUsage: jest.fn().mockReturnValue({
+      promptTokens: 100,
+      completionTokens: 50,
+      totalTokens: 150,
+    }),
+  },
 }));
 
 describe('QueryProcessor', () => {

packages/agents/__tests__/ragAgentFactory.test.ts

@@ -3,10 +3,10 @@ import { RagPipeline } from '../src/core/pipeline/ragPipeline';
 import { AvailableAgents, LLMConfig, DocumentSource } from '../src/types';
 import { Embeddings } from '@langchain/core/embeddings';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
-import { VectorStore } from '../src/db/vectorStore';
 import { mockDeep, MockProxy } from 'jest-mock-extended';
 import { BaseMessage } from '@langchain/core/messages';
 import EventEmitter from 'events';
+import { VectorStore } from '../src/db/postgresVectorStore';
 
 // Mock the agent configuration and RagPipeline
 jest.mock('../src/config/agent', () => ({
@@ -103,7 +103,7 @@ describe('RagAgentFactory', () => {
       // Assert
       expect(RagPipeline).toHaveBeenCalledTimes(1);
       expect(emitter).toBeInstanceOf(EventEmitter);
-      
+
       // Check streaming option is passed
       const executeSpy = (RagPipeline as jest.Mock).mock.results[0].value
         .execute;

packages/agents/__tests__/ragPipeline.test.ts

@@ -28,7 +28,20 @@ jest.mock('../src/utils/index', () => ({
     info: jest.fn(),
     debug: jest.fn(),
     error: jest.fn(),
-  }
+  },
+  TokenTracker: {
+    resetSessionCounters: jest.fn(),
+    getSessionTokenUsage: jest.fn().mockReturnValue({
+      promptTokens: 100,
+      responseTokens: 50,
+      totalTokens: 150,
+    }),
+    trackFullUsage: jest.fn().mockReturnValue({
+      promptTokens: 100,
+      completionTokens: 50,
+      totalTokens: 150,
+    }),
+  },
 }));
 
 describe('RagPipeline', () => {

packages/agents/sample.config.toml

@@ -1,16 +1,16 @@
 [API_KEYS]
-OPENAI = ""
-GROQ = ""
-ANTHROPIC = ""
-DEEPSEEK = ""
-GEMINI = ""
+OPENAI = "<your_openai_api_key>"
+GROQ = "<your_groq_api_key>"
+ANTHROPIC = "<your_anthropic_api_key>"
+DEEPSEEK = "<your_deepseek_api_key>"
+GEMINI = "<your_gemini_api_key>"
 
 [VECTOR_DB]
-POSTGRES_USER="cairocoder"
-POSTGRES_PASSWORD="cairocoder"
-POSTGRES_DB="cairocoder"
-POSTGRES_HOST="postgres"
-POSTGRES_PORT="5432"
+POSTGRES_USER = "cairocoder"
+POSTGRES_PASSWORD = "cairocoder"
+POSTGRES_DB = "cairocoder"
+POSTGRES_HOST = "postgres"
+POSTGRES_PORT = "5432"
 
 [GENERAL]
 PORT = 3_001
@@ -22,4 +22,4 @@ DEFAULT_CHAT_MODEL = "Gemini Flash 2.5"
 DEFAULT_FAST_CHAT_PROVIDER = "gemini"
 DEFAULT_FAST_CHAT_MODEL = "Gemini Flash 2.5"
 DEFAULT_EMBEDDING_PROVIDER = "openai"
-DEFAULT_EMBEDDING_MODEL = "Text embedding 3 large"
+DEFAULT_EMBEDDING_MODEL = "Text embedding 3 large"

packages/agents/src/config/agent.ts

@@ -1,25 +1,41 @@
-import { basicContractTemplate } from './templates/contractTemplate';
-import { cairoCoderPrompts } from './prompts';
-import { basicTestTemplate } from './templates/testTemplate';
 import { VectorStore } from '../db/postgresVectorStore';
-import { DocumentSource, RagSearchConfig } from '../types';
+import { RagSearchConfig } from '../types';
+import { getAgent, getDefaultAgent } from './agents';
 
+// Legacy function for backward compatibility
 export const getAgentConfig = (vectorStore: VectorStore): RagSearchConfig => {
+  const agent = getDefaultAgent();
   return {
-    name: 'Cairo Coder',
-    prompts: cairoCoderPrompts,
+    name: agent.name,
+    prompts: agent.prompts,
     vectorStore,
-    contractTemplate: basicContractTemplate,
-    testTemplate: basicTestTemplate,
-    maxSourceCount: 15,
-    similarityThreshold: 0.4,
-    sources: [
-      DocumentSource.CAIRO_BOOK,
-      DocumentSource.CAIRO_BY_EXAMPLE,
-      DocumentSource.STARKNET_FOUNDRY,
-      DocumentSource.CORELIB_DOCS,
-      DocumentSource.OPENZEPPELIN_DOCS,
-      DocumentSource.SCARB_DOCS,
-    ],
+    contractTemplate: agent.templates?.contract,
+    testTemplate: agent.templates?.test,
+    maxSourceCount: agent.parameters.maxSourceCount,
+    similarityThreshold: agent.parameters.similarityThreshold,
+    sources: agent.sources,
+  };
+};
+
+// Get agent configuration by ID
+export const getAgentConfigById = async (
+  agentId: string,
+  vectorStore: VectorStore,
+): Promise<RagSearchConfig> => {
+  const agentConfig = getAgent(agentId);
+
+  if (!agentConfig) {
+    throw new Error(`Agent configuration not found: ${agentId}`);
+  }
+
+  return {
+    name: agentConfig.name,
+    prompts: agentConfig.prompts,
+    vectorStore,
+    contractTemplate: agentConfig.templates?.contract,
+    testTemplate: agentConfig.templates?.test,
+    maxSourceCount: agentConfig.parameters.maxSourceCount,
+    similarityThreshold: agentConfig.parameters.similarityThreshold,
+    sources: agentConfig.sources,
   };
 };

packages/agents/src/config/agents.ts

@@ -0,0 +1,75 @@
+import { DocumentSource } from '../types';
+import { cairoCoderPrompts, scarbPrompts } from './prompts';
+import { basicContractTemplate } from './templates/contractTemplate';
+import { basicTestTemplate } from './templates/testTemplate';
+
+export interface AgentConfiguration {
+  id: string;
+  name: string;
+  description: string;
+  sources: DocumentSource[];
+  prompts: {
+    searchRetrieverPrompt: string;
+    searchResponsePrompt: string;
+    noSourceFoundPrompt?: string;
+  };
+  templates?: {
+    contract?: string;
+    test?: string;
+  };
+  parameters: {
+    maxSourceCount: number;
+    similarityThreshold: number;
+  };
+}
+
+// All pre-configured agents
+export const AGENTS: Record<string, AgentConfiguration> = {
+  'cairo-coder': {
+    id: 'cairo-coder',
+    name: 'Cairo Coder',
+    description:
+      'Default Cairo language and smart contract assistant with full documentation access',
+    sources: [
+      DocumentSource.CAIRO_BOOK,
+      DocumentSource.CAIRO_BY_EXAMPLE,
+      DocumentSource.STARKNET_FOUNDRY,
+      DocumentSource.CORELIB_DOCS,
+      DocumentSource.OPENZEPPELIN_DOCS,
+      DocumentSource.SCARB_DOCS,
+    ],
+    prompts: cairoCoderPrompts,
+    templates: {
+      contract: basicContractTemplate,
+      test: basicTestTemplate,
+    },
+    parameters: {
+      maxSourceCount: 15,
+      similarityThreshold: 0.4,
+    },
+  },
+  'scarb-assistant': {
+    id: 'scarb-assistant',
+    name: 'Scarb Assistant',
+    description: 'Specialized Scarb build tool assistance',
+    sources: [DocumentSource.SCARB_DOCS],
+    prompts: scarbPrompts,
+    parameters: {
+      maxSourceCount: 10,
+      similarityThreshold: 0.5,
+    },
+  },
+};
+
+// Helper functions
+export function getAgent(agentId: string): AgentConfiguration | undefined {
+  return AGENTS[agentId];
+}
+
+export function listAgents(): AgentConfiguration[] {
+  return Object.values(AGENTS);
+}
+
+export function getDefaultAgent(): AgentConfiguration {
+  return AGENTS['cairo-coder'];
+}

packages/agents/src/config/prompts/index.ts

@@ -1,20 +1,33 @@
-
 import {
   CAIROCODER_NO_SOURCE_PROMPT,
   CAIROCODER_RESPONSE_PROMPT,
   CAIROCODER_RETRIEVER_PROMPT,
 } from './cairoCoderPrompts';
 
-import { AgentPrompts } from '../../types';
+import {
+  SCARB_NO_SOURCE_PROMPT,
+  SCARB_RESPONSE_PROMPT,
+  SCARB_RETRIEVER_PROMPT,
+} from './scarbPrompts';
 
+import { AgentPrompts } from '../../types';
 
 export const cairoCoderPrompts: AgentPrompts = {
   searchRetrieverPrompt: CAIROCODER_RETRIEVER_PROMPT,
   searchResponsePrompt: CAIROCODER_RESPONSE_PROMPT,
   noSourceFoundPrompt: CAIROCODER_NO_SOURCE_PROMPT,
 };
 
+export const scarbPrompts: AgentPrompts = {
+  searchRetrieverPrompt: SCARB_RETRIEVER_PROMPT,
+  searchResponsePrompt: SCARB_RESPONSE_PROMPT,
+  noSourceFoundPrompt: SCARB_NO_SOURCE_PROMPT,
+};
+
 // Helper function to inject dynamic values into prompts
 export const injectPromptVariables = (prompt: string): string => {
-  return prompt.replace('${new Date().toISOString()}', new Date().toISOString());
+  return prompt.replace(
+    '${new Date().toISOString()}',
+    new Date().toISOString(),
+  );
 };