Added Ollama token counts

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 \

requirements.txt

@@ -20,3 +20,4 @@ pyyaml
 prometheus-client
 pyarrow
 boto3
+ollama

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
     }
   ]
 }

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})