From 46fd2d05941dc56f416a656d5f04b518ad601034 Mon Sep 17 00:00:00 2001
From: Han Xiao <han.xiao@jina.ai>
Date: Sat, 1 Mar 2025 22:30:52 +0800
Subject: [PATCH] refactor: agent schema

---
 src/agent.ts           |  3 ++-
 src/tools/evaluator.ts |  5 +++--
 src/utils/schemas.ts   | 18 +++++++++---------
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/agent.ts b/src/agent.ts
index 068105b..aab78e1 100644
--- a/src/agent.ts
+++ b/src/agent.ts
@@ -258,7 +258,8 @@ export async function getResponse(question?: string,
     messages = [{role: 'user', content: question.trim()}]
   }
 
-  const SchemaGen = new Schemas(question);
+  const SchemaGen = new Schemas();
+  await SchemaGen.setLanguage(question)
   const context: TrackerContext = {
     tokenTracker: existingContext?.tokenTracker || new TokenTracker(tokenBudget),
     actionTracker: existingContext?.actionTracker || new ActionTracker()
diff --git a/src/tools/evaluator.ts b/src/tools/evaluator.ts
index 6b8562a..71f82b8 100644
--- a/src/tools/evaluator.ts
+++ b/src/tools/evaluator.ts
@@ -366,7 +366,7 @@ Answer: ${answer}`
 
 function getQuestionEvaluationPrompt(question: string): PromptPair {
   return {
-    system: `You are an evaluator that determines if a question requires freshness, plurality, and/or completeness checks in addition to the required definitiveness check.
+    system: `You are an evaluator that determines if a question requires freshness, plurality, and/or completeness checks.
 
 <evaluation_types>
 1. freshness - Checks if the question is time-sensitive or requires very recent information
@@ -399,7 +399,7 @@ function getQuestionEvaluationPrompt(question: string): PromptPair {
      * Named time periods: "Renaissance and Industrial Revolution"
    - Look for explicitly named elements separated by commas, "and", "or", bullets
    - Example patterns: "comparing X and Y", "differences between A, B, and C", "both P and Q"
-   - DO NOT trigger for elements that aren't specifically named
+   - DO NOT trigger for elements that aren't specifically named   
 </rules>
 
 <examples>
@@ -498,6 +498,7 @@ Hier geht's um Investieren in der 'heutigen Wirtschaft', also brauche ich aktuel
 </output>
 </example-8>
 </examples>
+
 `,
     user:
 `${question}
diff --git a/src/utils/schemas.ts b/src/utils/schemas.ts
index 2b4ef71..dc0f786 100644
--- a/src/utils/schemas.ts
+++ b/src/utils/schemas.ts
@@ -64,18 +64,18 @@ export class Schemas {
   public languageCode: string = 'en';
 
 
-  constructor(query: string) {
+  async setLanguage(query: string) {
     const generator = new ObjectGeneratorSafe();
 
-    generator.generateObject({
+    const result = await generator.generateObject({
       model: 'evaluator',
       schema: this.getLanguageSchema(),
       prompt: getLanguagePrompt(query.slice(0, 100)),
-    }).then((result) => {
-      this.languageCode = result.object.langCode;
-      this.languageStyle = result.object.langStyle;
-      console.log(`langauge`, result.object);
     });
+
+    this.languageCode = result.object.langCode;
+    this.languageStyle = result.object.langStyle;
+    console.log(`langauge`, result.object);
   }
 
   getLanguagePrompt() {
@@ -91,10 +91,10 @@ export class Schemas {
 
   getQuestionEvaluateSchema(): z.ZodObject<any> {
     return z.object({
-      needsFreshness: z.boolean().describe('If the question requires freshness check'),
-      needsPlurality: z.boolean().describe('If the question requires plurality check'),
-      needsCompleteness: z.boolean().describe('If the question requires completeness check'),
       think: z.string().describe(`A very concise explain of why those checks are needed. ${this.getLanguagePrompt()}`).max(500),
+      needsFreshness: z.boolean(),
+      needsPlurality: z.boolean(),
+      needsCompleteness: z.boolean(),
     });
   }