Skip to content

Commit

Permalink
fix: streaming
Browse files Browse the repository at this point in the history
  • Loading branch information
hanxiao committed Feb 22, 2025
1 parent 86fb712 commit 7ec22a2
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 11 deletions.
10 changes: 6 additions & 4 deletions src/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -297,7 +297,7 @@ export async function getResponse(question?: string,
maxBadAttempts: number = 3,
existingContext?: Partial<TrackerContext>,
messages?: Array<CoreAssistantMessage | CoreUserMessage>
): Promise<{ result: StepAction; context: TrackerContext }> {
): Promise<{ result: StepAction; context: TrackerContext; visitedURLs: string[] }> {
const context: TrackerContext = {
tokenTracker: existingContext?.tokenTracker || new TokenTracker(tokenBudget),
actionTracker: existingContext?.actionTracker || new ActionTracker()
Expand Down Expand Up @@ -799,7 +799,7 @@ But unfortunately, you failed to solve the issue. You need to think out of the b
console.log(thisStep)

await storeContext(system, schema, [allContext, allKeywords, allQuestions, allKnowledge], totalStep);
return {result: thisStep, context};
return {result: thisStep, context, visitedURLs};

}

Expand Down Expand Up @@ -841,9 +841,11 @@ export async function main() {
const question = process.argv[2] || "";
const {
result: finalStep,
context: tracker
} = await getResponse(question) as { result: AnswerAction; context: TrackerContext };
context: tracker,
visitedURLs: visitedURLs
} = await getResponse(question) as { result: AnswerAction; context: TrackerContext; visitedURLs: string[] };
console.log('Final Answer:', finalStep.answer);
console.log('Visited URLs:', visitedURLs);

tracker.tokenTracker.printSummary();
}
Expand Down
15 changes: 9 additions & 6 deletions src/app.ts
Original file line number Diff line number Diff line change
Expand Up @@ -559,7 +559,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
}

try {
const {result: finalStep} = await getResponse(undefined, tokenBudget, maxBadAttempts, context, body.messages)
const {result: finalStep, visitedURLs: visitedURLs} = await getResponse(undefined, tokenBudget, maxBadAttempts, context, body.messages)

const usage = context.tokenTracker.getTotalUsageSnakeCase();
if (body.stream) {
Expand Down Expand Up @@ -595,7 +595,8 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
logprobs: null,
finish_reason: 'stop'
}],
usage
usage,
visitedURLs
};
res.write(`data: ${JSON.stringify(finalChunk)}\n\n`);
res.end();
Expand All @@ -616,15 +617,17 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
logprobs: null,
finish_reason: 'stop'
}],
usage
usage,
visitedURLs
};

// Log final response (excluding full content for brevity)
console.log('[chat/completions] Response:', {
id: response.id,
status: 200,
contentLength: response.choices[0].message.content.length,
usage: response.usage
usage: response.usage,
visitedURLs: response.visitedURLs
});

res.json(response);
Expand Down Expand Up @@ -662,7 +665,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
logprobs: null,
finish_reason: null
}],
usage
usage,
};
res.write(`data: ${JSON.stringify(closeThinkChunk)}\n\n`);

Expand Down Expand Up @@ -700,7 +703,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
logprobs: null,
finish_reason: 'stop'
}],
usage
usage,
};
res.json(response);
}
Expand Down
2 changes: 1 addition & 1 deletion src/tools/jina-dedup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import {TokenTracker} from "../utils/token-tracker";
import {JINA_API_KEY} from "../config";

const JINA_API_URL = 'https://api.jina.ai/v1/embeddings';
const SIMILARITY_THRESHOLD = 0.88; // Adjustable threshold for cosine similarity
const SIMILARITY_THRESHOLD = 0.85; // Adjustable threshold for cosine similarity

const JINA_API_CONFIG = {
MODEL: 'jina-embeddings-v3',
Expand Down
2 changes: 2 additions & 0 deletions src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,7 @@ export interface ChatCompletionResponse {
completion_tokens: number;
total_tokens: number;
};
visitedURLs?: string[];
}

export interface ChatCompletionChunk {
Expand All @@ -249,6 +250,7 @@ export interface ChatCompletionChunk {
finish_reason: null | 'stop';
}>;
usage?: any;
visitedURLs?: string[];
}

// Tracker Types
Expand Down

0 comments on commit 7ec22a2

Please sign in to comment.