From d719d97521a554b5e4361db5c4fcc35aae1e07cf Mon Sep 17 00:00:00 2001
From: JackColquitt <daniel@kalntera.ai>
Date: Fri, 27 Sep 2024 20:35:19 -0700
Subject: [PATCH] Added Ollama token counts

---
 Containerfile                                 |  2 +-
 requirements.txt                              |  1 +
 trustgraph/metering/pricelist.py              |  5 +++++
 .../model/text_completion/ollama/llm.py       | 20 +++++++++++++++----
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/Containerfile b/Containerfile
index 1bb0a7fd..7f80a514 100644
--- a/Containerfile
+++ b/Containerfile
@@ -13,7 +13,7 @@ RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \
 
 RUN pip3 install torch --index-url https://download.pytorch.org/whl/cpu
 
-RUN pip3 install anthropic boto3 cohere openai google-cloud-aiplatform \
+RUN pip3 install anthropic boto3 cohere openai google-cloud-aiplatform ollama \
     langchain langchain-core langchain-huggingface langchain-text-splitters \
     langchain-community pymilvus sentence-transformers transformers \
     huggingface-hub pulsar-client cassandra-driver pyarrow pyyaml \
diff --git a/requirements.txt b/requirements.txt
index 9a49a5aa..0d269066 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -20,3 +20,4 @@ pyyaml
 prometheus-client
 pyarrow
 boto3
+ollama
diff --git a/trustgraph/metering/pricelist.py b/trustgraph/metering/pricelist.py
index 6ded5e7d..bffa9bec 100644
--- a/trustgraph/metering/pricelist.py
+++ b/trustgraph/metering/pricelist.py
@@ -39,6 +39,11 @@ price_list = {
       "model_name": "cohere.command-r-plus-v1:0",
       "input_price": 0.0000030,
       "output_price": 0.0000150
+    },
+    {
+      "model_name": "ollama",
+      "input_price": 0,
+      "output_price": 0
     }
   ]
 }
\ No newline at end of file
diff --git a/trustgraph/model/text_completion/ollama/llm.py b/trustgraph/model/text_completion/ollama/llm.py
index 93d89720..b506b3cd 100755
--- a/trustgraph/model/text_completion/ollama/llm.py
+++ b/trustgraph/model/text_completion/ollama/llm.py
@@ -4,7 +4,7 @@ Simple LLM service, performs text prompt completion using an Ollama service.
 Input is prompt, output is response.
 """
 
-from langchain_community.llms import Ollama
+from ollama import Client
 from prometheus_client import Histogram, Info
 
 from .... schema import TextCompletionRequest, TextCompletionResponse, Error
@@ -67,7 +67,8 @@ class Processor(ConsumerProducer):
             "ollama": ollama,
         })
 
-        self.llm = Ollama(base_url=ollama, model=model)
+        self.model = model
+        self.llm = Client(host=ollama)
 
     def handle(self, msg):
 
@@ -83,11 +84,16 @@ class Processor(ConsumerProducer):
         try:
 
             with __class__.text_completion_metric.time():
-                response = self.llm.invoke(prompt)
+                response = self.llm.generate(self.model, prompt)
 
+            response_text = response['response']
             print("Send response...", flush=True)
+            print(response_text, flush=True)
 
-            r = TextCompletionResponse(response=response, error=None)
+            inputtokens = int(response['prompt_eval_count'])
+            outputtokens = int(response['eval_count'])
+
+            r = TextCompletionResponse(response=response_text, error=None, in_token=inputtokens, out_token=outputtokens, model="ollama")
 
             self.send(r, properties={"id": id})
 
@@ -105,6 +111,9 @@ class Processor(ConsumerProducer):
                     message = str(e),
                 ),
                 response=None,
+                in_token=None,
+                out_token=None,
+                model=None,
             )
 
             self.producer.send(r, properties={"id": id})
@@ -123,6 +132,9 @@ class Processor(ConsumerProducer):
                     message = str(e),
                 ),
                 response=None,
+                in_token=None,
+                out_token=None,
+                model=None,
             )
 
             self.producer.send(r, properties={"id": id})