search-query-builder.ts 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. import { parseJsonResponse } from '@/lib/generation/json-repair';
  2. import { PROMPT_IDS, buildPrompt } from '@/lib/generation/prompts';
  3. import type { AICallFn } from '@/lib/generation/pipeline-types';
  4. import { createLogger } from '@/lib/logger';
  5. const log = createLogger('SearchQueryBuilder');
  6. const TAVILY_SOFT_MAX_QUERY_LENGTH = 350;
  7. export const SEARCH_QUERY_REWRITE_EXCERPT_LENGTH = 7000;
  8. interface SearchQueryRewriteResponse {
  9. query: string;
  10. }
  11. export interface SearchQueryBuildResult {
  12. query: string;
  13. rewriteAttempted: boolean;
  14. rawRequirementLength: number;
  15. finalQueryLength: number;
  16. hasPdfContext: boolean;
  17. }
  18. function normalizeSearchRequirement(requirement: string): string {
  19. return requirement.replace(/\s+/g, ' ').trim();
  20. }
  21. function normalizePdfExcerpt(pdfText?: string): string {
  22. if (!pdfText) {
  23. return '';
  24. }
  25. return pdfText.replace(/\s+/g, ' ').trim().slice(0, SEARCH_QUERY_REWRITE_EXCERPT_LENGTH);
  26. }
  27. function shouldRewriteSearchQuery(
  28. normalizedRequirement: string,
  29. normalizedPdfExcerpt: string,
  30. ): boolean {
  31. return normalizedRequirement.length > 400 || Boolean(normalizedPdfExcerpt);
  32. }
  33. export async function buildSearchQuery(
  34. requirement: string,
  35. pdfText: string | undefined,
  36. aiCall?: AICallFn,
  37. ): Promise<SearchQueryBuildResult> {
  38. const normalizedRequirement = normalizeSearchRequirement(requirement);
  39. const pdfExcerpt = normalizePdfExcerpt(pdfText);
  40. const hasPdfContext = Boolean(pdfExcerpt);
  41. const rewriteAttempted = shouldRewriteSearchQuery(normalizedRequirement, pdfExcerpt);
  42. const fallback = {
  43. query: normalizedRequirement,
  44. rewriteAttempted,
  45. rawRequirementLength: normalizedRequirement.length,
  46. finalQueryLength: normalizedRequirement.length,
  47. hasPdfContext,
  48. } satisfies SearchQueryBuildResult;
  49. if (!normalizedRequirement || !rewriteAttempted) {
  50. return fallback;
  51. }
  52. if (!aiCall) {
  53. log.warn('Query rewrite AI call unavailable, falling back to raw requirement');
  54. return fallback;
  55. }
  56. const prompts = buildPrompt(PROMPT_IDS.WEB_SEARCH_QUERY_REWRITE, {
  57. requirement: normalizedRequirement,
  58. pdfExcerpt: pdfExcerpt || 'None',
  59. });
  60. if (!prompts) {
  61. log.warn('Query rewrite prompt not found, falling back to raw requirement');
  62. return fallback;
  63. }
  64. try {
  65. const response = await aiCall(prompts.system, prompts.user);
  66. const parsed = parseJsonResponse<SearchQueryRewriteResponse>(response);
  67. const rewrittenQuery = normalizeSearchRequirement(parsed?.query || '').slice(
  68. 0,
  69. TAVILY_SOFT_MAX_QUERY_LENGTH,
  70. );
  71. if (!rewrittenQuery) {
  72. log.warn('Query rewrite returned empty output, falling back to raw requirement');
  73. return fallback;
  74. }
  75. return {
  76. ...fallback,
  77. query: rewrittenQuery,
  78. finalQueryLength: rewrittenQuery.length,
  79. };
  80. } catch (error) {
  81. log.warn('Query rewrite failed, falling back to raw requirement:', error);
  82. return fallback;
  83. }
  84. }