refactor: consolidate type definitions into types.ts (#5)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Han Xiao <han.xiao@jina.ai>

src/agent.ts

@@ -9,7 +9,7 @@ import {evaluateAnswer} from "./tools/evaluator";
 import {analyzeSteps} from "./tools/error-analyzer";
 import {GEMINI_API_KEY, JINA_API_KEY, MODEL_NAME, SEARCH_PROVIDER, STEP_SLEEP} from "./config";
 import {tokenTracker} from "./utils/token-tracker";
-import {StepAction} from "./types";
+import {StepAction, SchemaProperty, ResponseSchema} from "./types";
 
 async function sleep(ms: number) {
   const seconds = Math.ceil(ms / 1000);
@@ -17,27 +17,6 @@ async function sleep(ms: number) {
   return new Promise(resolve => setTimeout(resolve, ms));
 }
 
-type SchemaProperty = {
-  type: SchemaType;
-  description: string;
-  enum?: string[];
-  items?: {
-    type: SchemaType;
-    description?: string;
-    properties?: Record<string, SchemaProperty>;
-    required?: string[];
-  };
-  properties?: Record<string, SchemaProperty>;
-  required?: string[];
-  maxItems?: number;
-};
-
-type ResponseSchema = {
-  type: SchemaType;
-  properties: Record<string, SchemaProperty>;
-  required: string[];
-};
-
 function getSchema(allowReflect: boolean, allowRead: boolean, allowAnswer: boolean, allowSearch: boolean): ResponseSchema {
   const actions: string[] = [];
   const properties: Record<string, SchemaProperty> = {
@@ -672,4 +651,4 @@ export async function main() {
 
 if (require.main === module) {
   main().catch(console.error);
-}
+}

src/tools/brave-search.ts

@@ -1,15 +1,7 @@
 import axios from 'axios';
 import {BRAVE_API_KEY} from "../config";
 
-interface BraveSearchResponse {
-  web: {
-    results: Array<{
-      title: string;
-      description: string;
-      url: string;
-    }>;
-  };
-}
+import { BraveSearchResponse } from '../types';
 
 export async function braveSearch(query: string): Promise<{ response: BraveSearchResponse }> {
   const response = await axios.get<BraveSearchResponse>('https://api.search.brave.com/res/v1/web/search', {
@@ -27,4 +19,4 @@ export async function braveSearch(query: string): Promise<{ response: BraveSearc
 
   // Maintain the same return structure as the original code
   return { response: response.data };
-}
+}

src/tools/dedup.ts

@@ -2,10 +2,7 @@ import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
 import { GEMINI_API_KEY, MODEL_NAME } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 
-type DedupResponse = {
-  thought: string;
-  unique_queries: string[];
-};
+import { DedupResponse } from '../types';
 
 const responseSchema = {
   type: SchemaType.OBJECT,
@@ -143,4 +140,4 @@ export async function main() {
 
 if (require.main === module) {
   main().catch(console.error);
-}
+}

src/tools/error-analyzer.ts

@@ -2,11 +2,7 @@ import {GoogleGenerativeAI, SchemaType} from "@google/generative-ai";
 import { GEMINI_API_KEY, MODEL_NAME } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 
-type EvaluationResponse = {
-    recap: string;
-    blame: string;
-    improvement: string;
-};
+import { ErrorAnalysisResponse } from '../types';
 
 const responseSchema = {
   type: SchemaType.OBJECT,
@@ -117,13 +113,13 @@ ${diaryContext.join('\n')}
 `;
 }
 
-export async function analyzeSteps(diaryContext: string[]): Promise<{ response: EvaluationResponse, tokens: number }> {
+export async function analyzeSteps(diaryContext: string[]): Promise<{ response: ErrorAnalysisResponse, tokens: number }> {
   try {
     const prompt = getPrompt(diaryContext);
     const result = await model.generateContent(prompt);
     const response = await result.response;
     const usage = response.usageMetadata;
-    const json = JSON.parse(response.text()) as EvaluationResponse;
+    const json = JSON.parse(response.text()) as ErrorAnalysisResponse;
     console.log('Error analysis:', {
       is_valid: !json.blame,
       reason: json.blame || 'No issues found'

src/tools/evaluator.ts

@@ -2,10 +2,7 @@ import { GoogleGenerativeAI, SchemaType } from "@google/generative-ai";
 import { GEMINI_API_KEY, MODEL_NAME } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 
-type EvaluationResponse = {
-  is_definitive: boolean;
-  reasoning: string;
-};
+import { EvaluationResponse } from '../types';
 
 const responseSchema = {
   type: SchemaType.OBJECT,
@@ -105,4 +102,4 @@ async function main() {
 
 if (require.main === module) {
   main().catch(console.error);
-}
+}

src/tools/getURLIndex.ts

@@ -1,24 +1,6 @@
-interface SearchResult {
-  title: string;
-  url: string;
-  description: string;
-}
+import { SearchResult, QueryResult, StepData } from '../types';
 
-interface QueryResult {
-  query: string;
-  results: SearchResult[];
-}
-
-export interface StepData {
-  step: number;
-  question: string;
-  action: string;
-  reasoning: string;
-  searchQuery?: string;
-  result?: QueryResult[];
-}
-
-export function buildURLMap(data: StepData[]): Record<string, string> {
+export function buildURLMap(data: StepData[]): Record<string, SearchResult['url'], QueryResult['query']> {
   const urlMap: Record<string, string> = {};
 
   data.forEach(step => {

src/tools/query-rewriter.ts

@@ -3,10 +3,7 @@ import { GEMINI_API_KEY, MODEL_NAME } from "../config";
 import { tokenTracker } from "../utils/token-tracker";
 import { SearchAction } from "../types";
 
-type KeywordsResponse = {
-  thought: string;
-  queries: string[];
-};
+import { KeywordsResponse } from '../types';
 
 const responseSchema = {
   type: SchemaType.OBJECT,
@@ -125,4 +122,4 @@ export async function rewriteQuery(action: SearchAction): Promise<{ queries: str
     console.error('Error in query rewriting:', error);
     throw error;
   }
-}
+}

src/tools/read.ts

@@ -1,17 +1,7 @@
 import https from 'https';
 import { tokenTracker } from "../utils/token-tracker";
 
-interface ReadResponse {
-  code: number;
-  status: number;
-  data: {
-    title: string;
-    description: string;
-    url: string;
-    content: string;
-    usage: { tokens: number; };
-  };
-}
+import { ReadResponse } from '../types';
 
 export function readUrl(url: string, token: string): Promise<{ response: ReadResponse, tokens: number }> {
   return new Promise((resolve, reject) => {

src/tools/search.ts

@@ -1,17 +1,7 @@
 import https from 'https';
 import { tokenTracker } from "../utils/token-tracker";
 
-interface SearchResponse {
-  code: number;
-  status: number;
-  data: Array<{
-    title: string;
-    description: string;
-    url: string;
-    content: string;
-    usage: { tokens: number; };
-  }>;
-}
+import { SearchResponse } from '../types';
 
 export function search(query: string, token: string): Promise<{ response: SearchResponse, tokens: number }> {
   return new Promise((resolve, reject) => {

src/types.ts

@@ -1,3 +1,6 @@
+import { SchemaType } from "@google/generative-ai";
+
+// Action Types
 type BaseAction = {
   action: "search" | "answer" | "reflect" | "visit";
   thoughts: string;
@@ -27,4 +30,107 @@ export type VisitAction = BaseAction & {
   URLTargets: string[];
 };
 
-export type StepAction = SearchAction | AnswerAction | ReflectAction | VisitAction;
+export type StepAction = SearchAction | AnswerAction | ReflectAction | VisitAction;
+
+// Response Types
+export interface TokenUsage {
+  tool: string;
+  tokens: number;
+}
+
+export interface SearchResponse {
+  code: number;
+  status: number;
+  data: Array<{
+    title: string;
+    description: string;
+    url: string;
+    content: string;
+    usage: { tokens: number; };
+  }>;
+}
+
+export interface BraveSearchResponse {
+  web: {
+    results: Array<{
+      title: string;
+      description: string;
+      url: string;
+    }>;
+  };
+}
+
+export type DedupResponse = {
+  thought: string;
+  unique_queries: string[];
+};
+
+export interface ReadResponse {
+  code: number;
+  status: number;
+  data: {
+    title: string;
+    description: string;
+    url: string;
+    content: string;
+    usage: { tokens: number; };
+  };
+}
+
+export type EvaluationResponse = {
+  is_definitive: boolean;
+  reasoning: string;
+};
+
+export type ErrorAnalysisResponse = {
+  recap: string;
+  blame: string;
+  improvement: string;
+};
+
+export interface SearchResult {
+  title: string;
+  url: string;
+  description: string;
+}
+
+export interface QueryResult {
+  query: string;
+  results: SearchResult[];
+}
+
+export interface StepData {
+  step: number;
+  question: string;
+  action: string;
+  reasoning: string;
+  searchQuery?: string;
+  result?: QueryResult[];
+}
+
+export type KeywordsResponse = {
+  thought: string;
+  queries: string[];
+};
+
+// Schema Types
+export type SchemaProperty = {
+  type: SchemaType;
+  description: string;
+  enum?: string[];
+  items?: {
+    type: SchemaType;
+    description?: string;
+    properties?: Record<string, SchemaProperty>;
+    required?: string[];
+  };
+  properties?: Record<string, SchemaProperty>;
+  required?: string[];
+  maxItems?: number;
+};
+
+export type ResponseSchema = {
+  type: SchemaType;
+  properties: Record<string, SchemaProperty>;
+  required: string[];
+};

src/utils/token-tracker.ts

@@ -1,9 +1,6 @@
 import { EventEmitter } from 'events';
 
-interface TokenUsage {
-  tool: string;
-  tokens: number;
-}
+import { TokenUsage } from '../types';
 
 class TokenTracker extends EventEmitter {
   private usages: TokenUsage[] = [];