@lmnr-ai/lmnr 0.4.21 → 0.4.23

package/src/laminar.ts

@@ -3,7 +3,8 @@ import {
   PipelineRunRequest,
   EvaluationDatapoint,
   CreateEvaluationResponse,
-  GetDatapointsResponse
+  GetDatapointsResponse,
+  SemanticSearchResponse
 } from './types';
 import {
   Attributes,
@@ -228,6 +229,46 @@ export class Laminar {
     }
   }
 
+  /**
+   * Perform a semantic search on a dataset.
+   * 
+   * @param query - The query string to search with.
+   * @param datasetId - The ID of the dataset to search in.
+   * @param limit - The maximum number of results to return.
+   * @param threshold - The minimum score for the results to be returned.
+   * @returns Response object containing the search results in descending order of relevance.
+   */
+  public static async semanticSearch({
+    query,
+    datasetId,
+    limit,
+    threshold,
+  }: {
+    query: string,
+    datasetId: string,
+    limit?: number,
+    threshold?: number,
+  }): Promise<SemanticSearchResponse> {
+    const body = JSON.stringify({
+      query,
+      datasetId,
+      limit,
+      threshold,
+    });
+
+    const response = await fetch(`${this.baseHttpUrl}/v1/semantic-search`, {
+      method: 'POST',
+      headers: this.getHeaders(),
+      body,
+    });
+
+    if (!response.ok) {
+      throw new Error(`Failed to perform semantic search. Response: ${response.statusText}`);
+    }
+
+    return await response.json() as SemanticSearchResponse;
+  }
+
   /**
    * Associates an event with the current span. If event with such name never
    * existed, Laminar will create a new event and infer its type from the value.
@@ -269,51 +310,90 @@ export class Laminar {
 
   /**
    * Sets the session information for the current span and returns the
-   * context to use for the following spans.
+   * context to use for the following spans. Returns the result of the
+   * function execution, so can be used in an `await` statement.
    * 
    * @param sessionId - The session ID to associate with the context.
-   * @param userId - The user ID to associate with the context.
-   * @returns The updated context with the association properties.
+   * @param fn - Function to execute within the session context.
+   * @returns The result of the function execution.
    * 
    * @example
-   * import { context as contextApi } from '@opentelemetry/api';
-   * import { Laminar } from '@lmnr-ai/laminar';
-   * const context = Laminar.contextWithSession({ sessionId: "1234", userId: "5678" });
-   * contextApi.with(context, () => {
-   *    // Your code here
+   * import { Laminar, observe } from '@lmnr-ai/lmnr';
+   * const result = await Laminar.withSession("session1234", async () => {
+   *   // Your code here
    * });
    */
-  public static contextWithSession({
-    sessionId,
-    userId
-  }: { sessionId?: string, userId?: string }): Context {
+  public static withSession<T>(
+    sessionId: string,
+    fn: () => T,
+  ): T {
     const currentSpan = trace.getActiveSpan();
     if (currentSpan !== undefined && isSpanContextValid(currentSpan.spanContext())) {
       if (sessionId) {
-        currentSpan.setAttribute(SESSION_ID, sessionId);
-      }
-      if (userId) {
-        currentSpan.setAttribute(USER_ID, userId);
+        currentSpan.setAttribute(SESSION_ID, JSON.stringify(sessionId));
       }
     }
     let associationProperties = {};
     if (sessionId) {
-      associationProperties = { ...associationProperties, "session_id": sessionId };
-    }
-    if (userId) {
-      associationProperties = { ...associationProperties, "user_id": userId };
+      associationProperties = { ...associationProperties, "session_id": JSON.stringify(sessionId) };
     }
 
     let entityContext = contextApi.active();
     const currentAssociationProperties = entityContext.getValue(ASSOCIATION_PROPERTIES_KEY);
-    if (associationProperties) {
+    if (associationProperties && Object.keys(associationProperties).length > 0) {
       entityContext = entityContext.setValue(
         ASSOCIATION_PROPERTIES_KEY,
         { ...(currentAssociationProperties ?? {}), ...associationProperties },
       );
     }
+    return contextApi.with(entityContext, fn);
+  }
+
+  /**
+   * Sets the metadata for the current span and returns the context to use for
+   * the following spans. Returns the result of the function execution, so can
+   * be used in an `await` statement.
+   * 
+   * @param metadata - The metadata to associate with the context. Set of string key
+   * string value pairs.
+   * @param fn - Function to execute within the metadata context.
+   * @returns The result of the function execution.
+   * 
+   * @example
+   * import { Laminar } from '@lmnr-ai/lmnr';
+   * const result = await Laminar.withMetadata(
+   *   {
+   *     "my_metadata_key": "my_metadata_value"
+   *   },
+   *   async () => {
+   *     // Your code here
+   *   }
+   * );
+   */
+  public static withMetadata<T>(
+    metadata: Record<string, string>,
+    fn: () => T,
+  ): T {
+    const currentSpan = trace.getActiveSpan();
+    if (currentSpan !== undefined && isSpanContextValid(currentSpan.spanContext())) {
+      for (const [key, value] of Object.entries(metadata)) {
+        currentSpan.setAttribute(`${ASSOCIATION_PROPERTIES}.metadata.${key}`, JSON.stringify(value));
+      }
+    }
+    let metadataAttributes = {};
+    for (const [key, value] of Object.entries(metadata)) {
+      metadataAttributes = { ...metadataAttributes, [`metadata.${key}`]: JSON.stringify(value) };
+    }
 
-    return entityContext;
+    let entityContext = contextApi.active();
+    const currentAssociationProperties = entityContext.getValue(ASSOCIATION_PROPERTIES_KEY);
+    if (metadataAttributes && Object.keys(metadataAttributes).length > 0) {
+      entityContext = entityContext.setValue(
+        ASSOCIATION_PROPERTIES_KEY,
+        { ...(currentAssociationProperties ?? {}), ...metadataAttributes },
+      );
+    }
+    return contextApi.with(entityContext, fn);
   }
 
   /**

package/src/types.ts

@@ -62,6 +62,17 @@ export type GetDatapointsResponse<D, T> = {
   anyInProject: boolean;
 }
 
+export type SemanticSearchResult = {
+  datasetId: StringUUID;
+  data: Record<string, any>;
+  content: string;
+  score: number;
+}
+
+export type SemanticSearchResponse = {
+  results: SemanticSearchResult[];
+}
+
 /**
  * Span types to categorize spans.
  *