From 64bdb88a4b56e2038f58654e45ad4eca4feb033e Mon Sep 17 00:00:00 2001
From: Travis Fischer <fisch0920@gmail.com>
Date: Wed, 24 Jul 2024 03:02:42 -0700
Subject: [PATCH] feat: minor fixes for diffbot and parsing of structured
 output

---
 src/parse-structured-output.ts   | 13 +++++++----
 src/services/diffbot-client.ts   | 40 ++++----------------------------
 src/services/proxycurl-client.ts |  2 +-
 src/url-utils.ts                 | 10 +++++---
 4 files changed, 22 insertions(+), 43 deletions(-)

diff --git a/src/parse-structured-output.ts b/src/parse-structured-output.ts
index f93bd47..24e4db0 100644
--- a/src/parse-structured-output.ts
+++ b/src/parse-structured-output.ts
@@ -35,14 +35,17 @@ export function parseStructuredOutput<T>(
   outputSchema: ZodType<T>
 ): T {
   let result
-  if (outputSchema instanceof z.ZodArray) {
+  if (outputSchema instanceof z.ZodArray || 'element' in outputSchema) {
     result = parseArrayOutput(output)
-  } else if (outputSchema instanceof z.ZodObject) {
+  } else if (outputSchema instanceof z.ZodObject || 'omit' in outputSchema) {
     result = parseObjectOutput(output)
   } else if (outputSchema instanceof z.ZodBoolean) {
     result = parseBooleanOutput(output)
-  } else if (outputSchema instanceof z.ZodNumber) {
-    result = parseNumberOutput(output, outputSchema)
+  } else if (
+    outputSchema instanceof z.ZodNumber ||
+    'nonnegative' in outputSchema
+  ) {
+    result = parseNumberOutput(output, outputSchema as unknown as z.ZodNumber)
   } else {
     // Default to string output...
     result = output
@@ -69,6 +72,8 @@ export function safeParseStructuredOutput<T>(
       data
     }
   } catch (err: any) {
+    console.error(err)
+
     return {
       success: false,
       error: err.message
diff --git a/src/services/diffbot-client.ts b/src/services/diffbot-client.ts
index 3dee6cc..720669a 100644
--- a/src/services/diffbot-client.ts
+++ b/src/services/diffbot-client.ts
@@ -1,16 +1,9 @@
-import type { Simplify } from 'type-fest'
 import defaultKy, { type KyInstance } from 'ky'
 import pThrottle from 'p-throttle'
 import { z } from 'zod'
 
 import { aiFunction, AIFunctionsProvider } from '../fns.js'
-import {
-  assert,
-  getEnv,
-  omit,
-  sanitizeSearchParams,
-  throttleKy
-} from '../utils.js'
+import { assert, getEnv, sanitizeSearchParams, throttleKy } from '../utils.js'
 
 export namespace diffbot {
   export const API_BASE_URL = 'https://api.diffbot.com'
@@ -335,7 +328,7 @@ export namespace diffbot {
     hits: number
     kgversion: string
     request_ctx: RequestCtx
-    data: EnhanceEntityResponseDatum[]
+    data: EnhanceEntityResult[]
     errors: any[]
   }
 
@@ -353,7 +346,7 @@ export namespace diffbot {
     search: string
   }
 
-  export interface EnhanceEntityResponseDatum {
+  export interface EnhanceEntityResult {
     score: number
     esscore: number
     entity: Entity
@@ -608,27 +601,6 @@ export namespace diffbot {
     name: string
     type: string
   }
-
-  export function pruneEntity(entity: diffbot.Entity) {
-    return omit(
-      entity,
-      'allOriginHashes',
-      'locations',
-      'images',
-      'nationalities',
-      'awards',
-      'interests',
-      'suppliers',
-      'partnerships',
-      'industries',
-      'categories',
-      'technographics',
-      'employeeCategories',
-      'diffbotClassification'
-    )
-  }
-
-  export type PrunedEntity = Simplify<ReturnType<typeof pruneEntity>>
 }
 
 /**
@@ -721,8 +693,8 @@ export class DiffbotClient extends AIFunctionsProvider {
   })
   async enhanceEntity(
     opts: diffbot.EnhanceEntityOptions
-  ): Promise<diffbot.PrunedEntity[]> {
-    const res = await this.kyKnowledgeGraph
+  ): Promise<diffbot.EnhanceEntityResponse> {
+    return this.kyKnowledgeGraph
       .get('kg/v3/enhance', {
         searchParams: sanitizeSearchParams({
           ...opts,
@@ -730,8 +702,6 @@ export class DiffbotClient extends AIFunctionsProvider {
         })
       })
       .json<diffbot.EnhanceEntityResponse>()
-
-    return res.data.map((datum) => diffbot.pruneEntity(datum.entity))
   }
 
   async searchKnowledgeGraph(options: diffbot.KnowledgeGraphSearchOptions) {
diff --git a/src/services/proxycurl-client.ts b/src/services/proxycurl-client.ts
index 5d13e45..4c88e85 100644
--- a/src/services/proxycurl-client.ts
+++ b/src/services/proxycurl-client.ts
@@ -1880,7 +1880,7 @@ export namespace proxycurl {
   export type PersonProfile = z.infer<typeof PersonProfileSchema>
 
   export type ResolvedPersonProfile = {
-    profile: PersonProfile
+    profile?: PersonProfile
     url?: string
     name_similarity_score?: number
     company_similarity_score?: number
diff --git a/src/url-utils.ts b/src/url-utils.ts
index 2b8c2cf..b11f5d2 100644
--- a/src/url-utils.ts
+++ b/src/url-utils.ts
@@ -1,6 +1,6 @@
 import isRelativeUrlImpl from 'is-relative-url'
 import normalizeUrlImpl, {
-  type Options as NormalizeUrlOptions
+  type Options as NormalizeUrlImplOptions
 } from 'normalize-url'
 import QuickLRU from 'quick-lru'
 
@@ -39,9 +39,13 @@ export function isRelativeUrl(url: string): boolean {
   return isRelativeUrlImpl(url) && !url.startsWith('//')
 }
 
+export type NormalizeUrlOptions = NormalizeUrlImplOptions & {
+  allowSloppyUris?: boolean
+}
+
 export function normalizeUrl(
   url?: string,
-  options?: NormalizeUrlOptions
+  { allowSloppyUris = true, ...options }: NormalizeUrlOptions = {}
 ): string | undefined {
   let normalizedUrl: string | undefined
 
@@ -50,7 +54,7 @@ export function normalizeUrl(
   }
 
   if (isRelativeUrl(url)) {
-    if (!/^[./]/.test(url) && url.indexOf('.') > 0) {
+    if (allowSloppyUris && !/^[#./]/.test(url) && url.indexOf('.') > 0) {
       url = `https://${url}`
     } else {
       return undefined