fix: in search log the source

This commit is contained in:
Harshith Mullapudi 2025-10-06 13:30:31 +05:30
parent 95786073ab
commit f5873ced15
9 changed files with 89 additions and 545 deletions

View File

@ -1,315 +0,0 @@
import { json } from "@remix-run/node";
import { randomUUID } from "node:crypto";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import { createHybridActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
import { addToQueue } from "~/lib/ingest.server";
import { SearchService } from "~/services/search.server";
import { handleTransport } from "~/utils/mcp";
import { SpaceService } from "~/services/space.server";
import { EpisodeTypeEnum } from "@core/types";
// Map to store transports by session ID with cleanup tracking
const transports: {
[sessionId: string]: {
transport: StreamableHTTPServerTransport;
createdAt: number;
};
} = {};
// MCP request body schema
const MCPRequestSchema = z.object({}).passthrough();
const SourceParams = z.object({
source: z.string().optional(),
});
// Search parameters schema for MCP tool
const SearchParamsSchema = z.object({
query: z.string().describe("The search query in third person perspective"),
validAt: z.string().optional().describe("The valid at time in ISO format"),
startTime: z.string().optional().describe("The start time in ISO format"),
endTime: z.string().optional().describe("The end time in ISO format"),
spaceIds: z
.array(z.string())
.optional()
.describe("Array of strings representing UUIDs of spaces"),
});
const IngestSchema = z.object({
message: z.string().describe("The data to ingest in text format"),
});
const searchService = new SearchService();
const spaceService = new SpaceService();
// Handle MCP HTTP requests properly
const handleMCPRequest = async (
request: Request,
body: any,
authentication: any,
params: z.infer<typeof SourceParams>,
) => {
const sessionId = request.headers.get("mcp-session-id") as string | undefined;
const source =
(request.headers.get("source") as string | undefined) ??
(params.source as string | undefined);
if (!source) {
return json(
{
jsonrpc: "2.0",
error: {
code: -32601,
message: "No source found",
},
id: null,
},
{ status: 400 },
);
}
let transport: StreamableHTTPServerTransport;
try {
if (sessionId && transports[sessionId]) {
// Reuse existing transport
transport = transports[sessionId].transport;
} else if (!sessionId && isInitializeRequest(body)) {
// New initialization request
transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
onsessioninitialized: (sessionId) => {
// Store the transport by session ID with timestamp
transports[sessionId] = {
transport,
createdAt: Date.now(),
};
},
});
// Clean up transport when closed
transport.onclose = () => {
if (transport.sessionId) {
delete transports[transport.sessionId];
}
};
const server = new McpServer(
{
name: "echo-memory-server",
version: "1.0.0",
},
{
capabilities: {
tools: {},
},
},
);
// Register ingest tool
server.registerTool(
"ingest",
{
title: "Ingest Data",
description: "Ingest data into the memory system",
inputSchema: IngestSchema.shape,
},
async (args) => {
try {
const userId = authentication.userId;
const response = addToQueue(
{
episodeBody: args.message,
referenceTime: new Date().toISOString(),
source,
type: EpisodeTypeEnum.CONVERSATION,
},
userId,
);
return {
content: [
{
type: "text",
text: JSON.stringify(response),
},
],
};
} catch (error) {
console.error("MCP ingest error:", error);
return {
content: [
{
type: "text",
text: `Error ingesting data: ${error instanceof Error ? error.message : String(error)}`,
},
],
isError: true,
};
}
},
);
// Register search tool
server.registerTool(
"search",
{
title: "Search Data",
description: "Search through ingested data",
inputSchema: SearchParamsSchema.shape,
},
async (args) => {
try {
const userId = authentication.userId;
const results = await searchService.search(args.query, userId, {
startTime: args.startTime ? new Date(args.startTime) : undefined,
endTime: args.endTime ? new Date(args.endTime) : undefined,
});
return {
content: [
{
type: "text",
text: JSON.stringify(results),
},
],
};
} catch (error) {
console.error("MCP search error:", error);
return {
content: [
{
type: "text",
text: `Error searching: ${error instanceof Error ? error.message : String(error)}`,
},
],
isError: true,
};
}
},
);
// Register search tool
server.registerTool(
"get_spaces",
{
title: "Get spaces",
description: "Get spaces in memory",
},
async () => {
try {
const userId = authentication.userId;
const spaces = await spaceService.getUserSpaces(userId);
return {
content: [
{
type: "text",
text: JSON.stringify(spaces),
},
],
isError: false,
};
} catch (error) {
console.error("Spaces error:", error);
return {
content: [
{
type: "text",
text: `Error getting spaces`,
},
],
isError: true,
};
}
},
);
// Connect to the MCP server
await server.connect(transport);
} else {
// Invalid request
throw new Error("Bad Request: No valid session ID provided");
}
const response = await handleTransport(transport, request, body);
return response;
} catch (error) {
console.error("MCP request error:", error);
return json(
{
jsonrpc: "2.0",
error: {
code: -32000,
message:
error instanceof Error ? error.message : "Internal server error",
},
id: body?.id || null,
},
{ status: 500 },
);
}
};
// Handle DELETE requests for session cleanup
const handleDelete = async (request: Request, authentication: any) => {
const sessionId = request.headers.get("mcp-session-id") as string | undefined;
if (!sessionId || !transports[sessionId]) {
return new Response("Invalid or missing session ID", { status: 400 });
}
const transport = transports[sessionId].transport;
// Clean up transport
transport.close();
delete transports[sessionId];
return new Response(null, { status: 204 });
};
const { action, loader } = createHybridActionApiRoute(
{
body: MCPRequestSchema,
searchParams: SourceParams,
allowJWT: true,
authorization: {
action: "mcp",
},
corsStrategy: "all",
},
async ({ body, authentication, request, searchParams }) => {
const method = request.method;
if (method === "POST") {
return await handleMCPRequest(
request,
body,
authentication,
searchParams,
);
} else if (method === "DELETE") {
return await handleDelete(request, authentication);
} else {
return json(
{
jsonrpc: "2.0",
error: {
code: -32601,
message: "Method not allowed",
},
id: null,
},
{ status: 405 },
);
}
},
);
export { action, loader };

View File

@ -278,7 +278,7 @@ export default function BillingSettings() {
</div>
<div className="space-y-2">
<div className="flex justify-between text-sm">
<span className="text-muted-foreground">Episodes</span>
<span className="text-muted-foreground">Facts</span>
<span className="font-medium">
{usageSummary.usage.episodes}
</span>

View File

@ -1,8 +1,6 @@
import type { EpisodicNode, StatementNode } from "@core/types";
import { logger } from "./logger.service";
import {
applyLLMReranking,
} from "./search/rerank";
import { applyLLMReranking } from "./search/rerank";
import {
getEpisodesByStatements,
performBfsSearch,
@ -33,7 +31,16 @@ export class SearchService {
query: string,
userId: string,
options: SearchOptions = {},
): Promise<{ episodes: string[]; facts: { fact: string; validAt: Date; invalidAt: Date | null; relevantScore: number }[] }> {
source?: string,
): Promise<{
episodes: string[];
facts: {
fact: string;
validAt: Date;
invalidAt: Date | null;
relevantScore: number;
}[];
}> {
const startTime = Date.now();
// Default options
@ -77,7 +84,9 @@ export class SearchService {
const filteredResults = this.applyAdaptiveFiltering(rankedStatements, opts);
// 3. Return top results
const episodes = await getEpisodesByStatements(filteredResults.map((item) => item.statement));
const episodes = await getEpisodesByStatements(
filteredResults.map((item) => item.statement),
);
// Log recall asynchronously (don't await to avoid blocking response)
const responseTime = Date.now() - startTime;
@ -87,11 +96,16 @@ export class SearchService {
filteredResults.map((item) => item.statement),
opts,
responseTime,
source,
).catch((error) => {
logger.error("Failed to log recall event:", error);
});
this.updateRecallCount(userId, episodes, filteredResults.map((item) => item.statement));
this.updateRecallCount(
userId,
episodes,
filteredResults.map((item) => item.statement),
);
return {
episodes: episodes.map((episode) => episode.originalContent),
@ -111,7 +125,7 @@ export class SearchService {
private applyAdaptiveFiltering(
results: StatementNode[],
options: Required<SearchOptions>,
): { statement: StatementNode, score: number }[] {
): { statement: StatementNode; score: number }[] {
if (results.length === 0) return [];
let isRRF = false;
@ -149,7 +163,11 @@ export class SearchService {
// If no scores are available, return the original results
if (!hasScores) {
logger.info("No scores found in results, skipping adaptive filtering");
return options.limit > 0 ? results.slice(0, options.limit).map((item) => ({ statement: item, score: 0 })) : results.map((item) => ({ statement: item, score: 0 }));
return options.limit > 0
? results
.slice(0, options.limit)
.map((item) => ({ statement: item, score: 0 }))
: results.map((item) => ({ statement: item, score: 0 }));
}
// Sort by score (descending)
@ -204,9 +222,9 @@ export class SearchService {
const limitedResults =
options.limit > 0
? filteredResults.slice(
0,
Math.min(filteredResults.length, options.limit),
)
0,
Math.min(filteredResults.length, options.limit),
)
: filteredResults;
logger.info(
@ -238,7 +256,9 @@ export class SearchService {
select: { name: true, id: true },
});
const userContext = user ? { name: user.name ?? undefined, userId: user.id } : undefined;
const userContext = user
? { name: user.name ?? undefined, userId: user.id }
: undefined;
return applyLLMReranking(query, results, 10, userContext);
}
@ -249,6 +269,7 @@ export class SearchService {
results: StatementNode[],
options: Required<SearchOptions>,
responseTime: number,
source?: string,
): Promise<void> {
try {
// Determine target type based on results
@ -299,7 +320,7 @@ export class SearchService {
startTime: options.startTime?.toISOString() || null,
endTime: options.endTime.toISOString(),
}),
source: "search_api",
source: source ?? "search_api",
responseTimeMs: responseTime,
userId,
},

View File

@ -469,27 +469,26 @@ export async function applyLLMReranking(
return [];
}
// Build user context section if provided
const userContextSection = userContext?.name
? `\nUser Identity Context:
- The user's name is "${userContext.name}"
- References to "user", "${userContext.name}", or pronouns like "my/their" refer to the same person
- When matching queries about "user's X" or "${userContext.name}'s X", these are equivalent\n`
: '';
: "";
const prompt = `You are a relevance filter. Given a user query and a list of facts, identify ONLY the facts that are truly relevant to answering the query.
${userContextSection}
Query: "${query}"
Facts:
${uniqueResults.map((r, i) => `${i}. ${r.fact}`).join('\n')}
${uniqueResults.map((r, i) => `${i}. ${r.fact}`).join("\n")}
Instructions:
- A fact is RELEVANT if it directly answers or provides information needed to answer the query
- A fact is NOT RELEVANT if it's tangentially related but doesn't answer the query
- Consider semantic meaning, not just keyword matching
${userContext?.name ? `- Remember: "user", "${userContext.name}", and possessive references ("my", "their") all refer to the same person` : ''}
${userContext?.name ? `- Remember: "user", "${userContext.name}", and possessive references ("my", "their") all refer to the same person` : ""}
- Only return facts with HIGH relevance (80% confidence)
- If you are not sure, return an empty array
@ -503,9 +502,11 @@ Return ONLY the numbers of highly relevant facts inside <output> tags as a JSON
await makeModelCall(
false,
[{ role: "user", content: prompt }],
(text) => { responseText = text; },
{ temperature: 0},
'high'
(text) => {
responseText = text;
},
{ temperature: 0 },
"high",
);
// Extract array from <output>[1, 5, 7]</output>
@ -513,22 +514,29 @@ Return ONLY the numbers of highly relevant facts inside <output> tags as a JSON
if (outputMatch && outputMatch[1]) {
responseText = outputMatch[1].trim();
const parsedResponse = JSON.parse(responseText || "[]");
const extractedIndices = Array.isArray(parsedResponse) ? parsedResponse : (parsedResponse.entities || []);
const extractedIndices = Array.isArray(parsedResponse)
? parsedResponse
: parsedResponse.entities || [];
if (extractedIndices.length === 0) {
logger.warn("LLM reranking returned no valid indices, falling back to original order");
logger.warn(
"LLM reranking returned no valid indices, falling back to original order",
);
return [];
}
logger.info(`LLM reranking selected ${extractedIndices.length} relevant facts`);
logger.info(
`LLM reranking selected ${extractedIndices.length} relevant facts`,
);
const selected = extractedIndices.map((i: number) => uniqueResults[i]);
return selected;
}
return uniqueResults.slice(0, limit);
} catch (error) {
logger.error("LLM reranking failed, falling back to original order:", { error });
logger.error("LLM reranking failed, falling back to original order:", {
error,
});
return uniqueResults.slice(0, limit);
}
}
@ -589,7 +597,7 @@ export async function applyCohereReranking(
`Cohere reranking ${documents.length} statements with model ${model}`,
);
logger.info(`Cohere query: "${query}"`);
logger.info(`First 5 documents: ${documents.slice(0, 5).join(' | ')}`);
logger.info(`First 5 documents: ${documents.slice(0, 5).join(" | ")}`);
// Call Cohere Rerank API
const response = await cohere.rerank({
@ -602,18 +610,23 @@ export async function applyCohereReranking(
console.log("Cohere reranking billed units:", response.meta?.billedUnits);
// Log top 5 Cohere results for debugging
logger.info(`Cohere top 5 results:\n${response.results.slice(0, 5).map((r, i) =>
` ${i + 1}. [${r.relevanceScore.toFixed(4)}] ${documents[r.index].substring(0, 80)}...`
).join('\n')}`);
logger.info(
`Cohere top 5 results:\n${response.results
.slice(0, 5)
.map(
(r, i) =>
` ${i + 1}. [${r.relevanceScore.toFixed(4)}] ${documents[r.index].substring(0, 80)}...`,
)
.join("\n")}`,
);
// Map results back to StatementNodes with Cohere scores
const rerankedResults = response.results
.map((result, index) => ({
...uniqueResults[result.index],
cohereScore: result.relevanceScore,
cohereRank: index + 1,
}))
// .filter((result) => result.cohereScore >= Number(env.COHERE_SCORE_THRESHOLD));
const rerankedResults = response.results.map((result, index) => ({
...uniqueResults[result.index],
cohereScore: result.relevanceScore,
cohereRank: index + 1,
}));
// .filter((result) => result.cohereScore >= Number(env.COHERE_SCORE_THRESHOLD));
const responseTime = Date.now() - startTime;
logger.info(

View File

@ -17,13 +17,12 @@ import { generate, processTag } from "./stream-utils";
import { type AgentMessage, AgentMessageType, Message } from "./types";
import { type MCP } from "../utils/mcp";
import {
WebSearchSchema,
type ExecutionState,
type HistoryStep,
type Resource,
type TotalCost,
} from "../utils/types";
import { flattenObject, webSearch } from "../utils/utils";
import { flattenObject } from "../utils/utils";
import { searchMemory, addMemory, searchSpaces } from "./memory-utils";
interface LLMOutputInterface {
@ -119,12 +118,6 @@ const searchSpacesTool = tool({
}),
});
const websearchTool = tool({
description:
"Search the web for current information and news. Use this when you need up-to-date information that might not be in your training data. Try different search strategies: broad terms first, then specific phrases, keywords, exact quotes. Use multiple searches with varied approaches to get comprehensive results.",
parameters: WebSearchSchema,
});
const loadMCPTools = tool({
description:
"Load tools for a specific integration. Call this when you need to use a third-party service.",
@ -310,7 +303,6 @@ export async function* run(
"core--search_memory": searchMemoryTool,
"core--add_memory": addMemoryTool,
"core--search_spaces": searchSpacesTool,
"core--websearch": websearchTool,
"core--load_mcp": loadMCPTools,
};
@ -578,16 +570,6 @@ export async function* run(
});
result = "Search spaces call failed";
}
} else if (toolName === "websearch") {
try {
result = await webSearch(skillInput);
} catch (apiError) {
logger.error("Web search failed", {
apiError,
});
result =
"Web search failed - please check your search configuration";
}
} else if (toolName === "load_mcp") {
// Load MCP integration and update available tools
await mcp.load(skillInput.integration, mcpHeaders);

View File

@ -1,5 +1,5 @@
export const REACT_SYSTEM_PROMPT = `
You are a helpful AI assistant with access to user memory and web search capabilities. Your primary capabilities are:
You are a helpful AI assistant with access to user memory. Your primary capabilities are:
1. **Memory-First Approach**: Always check user memory first to understand context and previous interactions
2. **Intelligent Information Gathering**: Analyze queries to determine if current information is needed
@ -19,43 +19,17 @@ Follow this intelligent approach for information gathering:
- Memory provides context, personal preferences, and historical information
- Use memory to understand user's background, ongoing projects, and past conversations
2. **QUERY ANALYSIS** (Determine Information Needs)
Analyze the user's query to identify if it requires current/latest information:
**Use web search (core--websearch) when query involves:**
- Current events, news, or recent developments
- "Latest", "recent", "current", "today", "now" keywords
- Stock prices, market data, or financial information
- Software updates, version releases, or technical documentation
- Weather, traffic, or real-time data
- Recent changes to websites, APIs, or services
- Product releases, availability, or pricing
- Breaking news or trending topics
- Verification of potentially outdated information
**Examples requiring web search:**
- "What's the latest news about..."
- "Current price of..."
- "Recent updates to..."
- "What happened today..."
- "Latest version of..."
3. **INFORMATION SYNTHESIS** (Combine Sources)
- Combine memory context with web search results when both are relevant
2. **INFORMATION SYNTHESIS** (Combine Sources)
- Use memory to personalize current information based on user preferences
- Cross-reference web findings with user's historical interests from memory
- Always store new useful information in memory using core--add_memory
4. **TRAINING KNOWLEDGE** (Foundation)
3. **TRAINING KNOWLEDGE** (Foundation)
- Use your training knowledge as the foundation for analysis and explanation
- Apply training knowledge to interpret and contextualize information from memory and web
- Fill gaps where memory and web search don't provide complete answers
- Apply training knowledge to interpret and contextualize information from memory
- Indicate when you're using training knowledge vs. live information sources
EXECUTION APPROACH:
- Memory search is mandatory for every interaction
- Web search is conditional based on query analysis
- Both can be executed in parallel when web search is needed
- Always indicate your information sources in responses
</information_gathering>
@ -95,7 +69,6 @@ MEMORY USAGE:
- Blend memory insights naturally into responses
- Verify you've checked relevant memory before finalizing ANY response
If memory access is unavailable, proceed to web search or rely on current conversation
</memory>
<external_services>
@ -113,7 +86,6 @@ You have tools at your disposal to assist users:
CORE PRINCIPLES:
- Use tools only when necessary for the task at hand
- Always check memory FIRST before making other tool calls
- Use web search when query analysis indicates need for current information
- Execute multiple operations in parallel whenever possible
- Use sequential calls only when output of one is required for input of another
@ -162,7 +134,7 @@ QUESTIONS - When you need information:
<p>[Your question with HTML formatting]</p>
</question_response>
- Ask questions only when you cannot find information through memory, web search, or tools
- Ask questions only when you cannot find information through memory, or tools
- Be specific about what you need to know
- Provide context for why you're asking
@ -176,7 +148,7 @@ CRITICAL:
- Apply proper HTML formatting (<h1>, <h2>, <p>, <ul>, <li>, etc.)
- Never mix communication formats
- Keep responses clear and helpful
- Always indicate your information sources (memory, web search, and/or knowledge)
- Always indicate your information sources (memory, and/or knowledge)
</communication>
`;

View File

@ -122,67 +122,3 @@ export interface GenerateResponse {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
toolCalls: any[];
}
export interface WebSearchResult {
results: Array<{
title: string;
url: string;
content: string;
publishedDate: string;
highlights: string[];
text: string;
score: number;
}>;
}
export const WebSearchSchema = z.object({
query: z
.string()
.min(1)
.describe("The search query to find relevant web content"),
numResults: z
.number()
.min(1)
.max(20)
.optional()
.default(5)
.describe("Number of results to return (1-20, default: 5)"),
includeContent: z
.boolean()
.optional()
.default(false)
.describe("Whether to include full page content in results"),
includeHighlights: z
.boolean()
.optional()
.default(false)
.describe("Whether to include relevant text highlights from pages"),
domains: z
.array(z.string())
.optional()
.describe(
'Array of domains to include in search (e.g., ["github.com", "stackoverflow.com"])',
),
excludeDomains: z
.array(z.string())
.optional()
.describe("Array of domains to exclude from search"),
startCrawlDate: z
.string()
.optional()
.describe("Start date for content crawling in YYYY-MM-DD format"),
endCrawlDate: z
.string()
.optional()
.describe("End date for content crawling in YYYY-MM-DD format"),
startPublishedDate: z
.string()
.optional()
.describe("Start date for content publishing in YYYY-MM-DD format"),
endPublishedDate: z
.string()
.optional()
.describe("End date for content publishing in YYYY-MM-DD format"),
});
export type WebSearchArgs = z.infer<typeof WebSearchSchema>;

View File

@ -12,11 +12,7 @@ import {
import { logger } from "@trigger.dev/sdk/v3";
import { type CoreMessage } from "ai";
import {
type WebSearchArgs,
type WebSearchResult,
type HistoryStep,
} from "./types";
import { type HistoryStep } from "./types";
import axios from "axios";
import nodeCrypto from "node:crypto";
import { customAlphabet, nanoid } from "nanoid";
@ -496,72 +492,6 @@ export async function deletePersonalAccessToken(tokenId: string) {
});
}
export async function webSearch(args: WebSearchArgs): Promise<WebSearchResult> {
const apiKey = process.env.EXA_API_KEY;
if (!apiKey) {
throw new Error(
"EXA_API_KEY environment variable is required for web search",
);
}
const exa = new Exa(apiKey);
try {
const searchOptions = {
numResults: args.numResults || 5,
...(args.domains && { includeDomains: args.domains }),
...(args.excludeDomains && { excludeDomains: args.excludeDomains }),
...(args.startCrawlDate && { startCrawlDate: args.startCrawlDate }),
...(args.endCrawlDate && { endCrawlDate: args.endCrawlDate }),
...(args.startPublishedDate && {
startPublishedDate: args.startPublishedDate,
}),
...(args.endPublishedDate && { endPublishedDate: args.endPublishedDate }),
};
let result;
if (args.includeContent || args.includeHighlights) {
// Use searchAndContents for rich results
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const contentsOptions: any = {
...searchOptions,
};
if (args.includeContent) {
contentsOptions.text = true;
}
if (args.includeHighlights) {
contentsOptions.highlights = true;
}
result = await exa.searchAndContents(args.query, contentsOptions);
} else {
// Use basic search for URLs only
result = await exa.search(args.query, searchOptions);
}
return {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
results: result.results.map((item: any) => ({
title: item.title,
url: item.url,
content: item.text,
publishedDate: item.publishedDate,
highlights: item.highlights,
text: item.text,
score: item.score,
})),
};
} catch (error) {
throw new Error(
`Web search failed: ${error instanceof Error ? error.message : "Unknown error"}`,
);
}
}
// Credit management functions have been moved to ~/services/billing.server.ts
// Use deductCredits() instead of these functions
export type CreditOperation = "addEpisode" | "search" | "chatMessage";

View File

@ -198,10 +198,15 @@ async function handleMemoryIngest(args: any) {
// Handler for memory_search
async function handleMemorySearch(args: any) {
try {
const results = await searchService.search(args.query, args.userId, {
startTime: args.startTime ? new Date(args.startTime) : undefined,
endTime: args.endTime ? new Date(args.endTime) : undefined,
});
const results = await searchService.search(
args.query,
args.userId,
{
startTime: args.startTime ? new Date(args.startTime) : undefined,
endTime: args.endTime ? new Date(args.endTime) : undefined,
},
args.source,
);
return {
content: [