From 77249ac88f2045bea61026c57abfdc336178ae93 Mon Sep 17 00:00:00 2001
From: Rohit Kadhe <rohit@scale3labs.com>
Date: Tue, 9 Jul 2024 09:48:05 -0600
Subject: [PATCH] add ollama and add more event types

---
 src/typescript/constants/common.ts | 12 ++++++++---
 src/typescript/constants/ollama.ts | 32 ++++++++++++++++++++++++++++++
 src/typescript/index.ts            |  2 ++
 src/typescript/package-lock.json   |  6 +++---
 src/typescript/package.json        |  2 +-
 5 files changed, 47 insertions(+), 7 deletions(-)
 create mode 100644 src/typescript/constants/ollama.ts

diff --git a/src/typescript/constants/common.ts b/src/typescript/constants/common.ts
index 178dbb6..a13f771 100644
--- a/src/typescript/constants/common.ts
+++ b/src/typescript/constants/common.ts
@@ -12,6 +12,7 @@ import { PineConeFunctionNames, PineConeFunctions } from './pinecone'
 import { QdrantFunctionNames, QdrantFunctions } from './qdrant'
 import { WeaviateFunctionNames, WeaviateFunctions } from './weaviate'
 import { TiktokenModel, TiktokenEncoding } from 'tiktoken'
+import { OllamaFunctionNames, OllamaFunctions } from './ollama'
 
 export const LLMSpanAttributeNames: Array<keyof typeof LLMSpanAttributesObj> = Object.keys(LLMSpanAttributesObj) as Array<keyof typeof LLMSpanAttributesObj>
 export const DatabaseSpanAttributeNames: Array<keyof typeof DatabaseSpanAttributesObj> = Object.keys(DatabaseSpanAttributesObj) as Array<keyof typeof DatabaseSpanAttributesObj>
@@ -28,7 +29,8 @@ export const Vendors = {
   CHROMADB: 'chromadb',
   QDRANT: 'qdrant',
   WEAVIATE: 'weaviate',
-  PG: 'pg'
+  PG: 'pg',
+  OLLAMA: 'ollama'
 } as const
 
 export enum Event {
@@ -36,6 +38,8 @@ export enum Event {
   STREAM_OUTPUT = 'stream.output',
   STREAM_END = 'stream.end',
   RESPONSE = 'response',
+  GEN_AI_COMPLETION = 'gen_ai.content.completion',
+  GEN_AI_PROMPT = 'gen_ai.content.prompt',
 }
 export type Vendor = typeof Vendors[keyof typeof Vendors]
 
@@ -49,7 +53,8 @@ interface VendorInstrumentationFunctions {
   chromadb: ChromadbFunctions[]
   qdrant: QdrantFunctions[]
   weaviate: WeaviateFunctions[]
-  pg: PgFunctions[]
+  pg: PgFunctions[],
+  ollama: OllamaFunctions[]
 }
 
 
@@ -68,7 +73,8 @@ export const TracedFunctionsByVendor: VendorTracedFunctions = {
   openai: OpenAIFunctionNames,
   pinecone: PineConeFunctionNames,
   qdrant: QdrantFunctionNames,
-  weaviate: WeaviateFunctionNames
+  weaviate: WeaviateFunctionNames,
+  ollama: OllamaFunctionNames
 } as const
 
 export const TIKTOKEN_MODEL_MAPPING: Record<TiktokenModel | string, TiktokenEncoding> = {
diff --git a/src/typescript/constants/ollama.ts b/src/typescript/constants/ollama.ts
new file mode 100644
index 0000000..1b8d780
--- /dev/null
+++ b/src/typescript/constants/ollama.ts
@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2024 Scale3 Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export const APIS = {
+  GENERATE: {
+    METHOD: 'ollama.generate',
+    ENDPOINT: '/api/generate'
+  },
+  CHAT: {
+    METHOD: 'ollama.chat',
+    ENDPOINT: '/api/chat'
+  },
+  EMBEDDINGS: {
+    METHOD: 'ollama.embeddings',
+    ENDPOINT: ' /api/embeddings'
+  }
+} as const
+export type OllamaFunctions = typeof APIS[keyof typeof APIS]['METHOD']
+export const OllamaFunctionNames: OllamaFunctions[] = Object.values(APIS).map((api) => api.METHOD)
\ No newline at end of file
diff --git a/src/typescript/index.ts b/src/typescript/index.ts
index 3ee825f..1a04b1e 100644
--- a/src/typescript/index.ts
+++ b/src/typescript/index.ts
@@ -11,6 +11,7 @@ import { APIS as LlamaIndexAPIS } from "./constants/llamaindex";
 import { APIS as OpenAIAPIs } from "./constants/openai";
 import { APIS as PineConeAPIS } from "./constants/pinecone";
 import { APIS as QdrantAPIS } from "./constants/qdrant";
+import { APIS as ollamaAPIS } from "./constants/ollama";
 import { queryTypeToFunctionToProps } from "./constants/weaviate";
 import { TIKTOKEN_MODEL_MAPPING } from "./constants/common";
 
@@ -45,6 +46,7 @@ const APIS = {
   openai: OpenAIAPIs,
   pinecone: PineConeAPIS,
   qdrant: QdrantAPIS,
+  ollama: ollamaAPIS
 }
 export {
   LLMSpanAttributeNames,
diff --git a/src/typescript/package-lock.json b/src/typescript/package-lock.json
index 609226e..d9c33ee 100644
--- a/src/typescript/package-lock.json
+++ b/src/typescript/package-lock.json
@@ -1,12 +1,12 @@
 {
     "name": "@langtrase/trace-attributes",
-    "version": "6.0.1",
+    "version": "6.0.2",
     "lockfileVersion": 3,
     "requires": true,
     "packages": {
         "": {
             "name": "@langtrase/trace-attributes",
-            "version": "6.0.1",
+            "version": "6.0.2",
             "license": "MIT",
             "dependencies": {
                 "json-schema-to-typescript": "^14.1.0",
@@ -856,4 +856,4 @@
             }
         }
     }
-}
+}
\ No newline at end of file
diff --git a/src/typescript/package.json b/src/typescript/package.json
index 538007c..269f89f 100644
--- a/src/typescript/package.json
+++ b/src/typescript/package.json
@@ -1,6 +1,6 @@
 {
     "name": "@langtrase/trace-attributes",
-    "version": "6.0.1",
+    "version": "6.0.2",
     "description": "LangTrace - Trace Attributes",
     "main": "dist/index.js",
     "types": "dist/index.d.ts",