better output messages.. this should really use logging

Author: deadbits

client.py

@@ -1,16 +1,31 @@
 #!/usr/bin/env python3
 # github.com/deadbits/vector-embedding-api
 # client.py
+import os
+import sys
 import json
 import argparse
 import requests
+from pydantic import BaseModel
+from typing import List, Optional
+from datetime import datetime
 
 
-def send_request(text, model_type='local'):
+def timestamp_str():
+    return datetime.isoformat(datetime.utcnow())
+
+
+class Embedding(BaseModel):
+    text: str = ''
+    embedding: List[float] = []
+    metadata: Optional[dict] = {}
+
+
+def send_request(text_batch, model_type='local'):
     url = 'http://127.0.0.1:5000/submit'
     headers = {'Content-Type': 'application/json'}
     payload = {
-        'text': text,
+        'text': text_batch,
         'model': model_type
     }
 
@@ -20,14 +35,33 @@ def send_request(text, model_type='local'):
             headers=headers,
             data=json.dumps(payload)
         )
-        
+
         response.raise_for_status()
         return response.json()
     except requests.RequestException as err:
         print(f'[error] exception sending http request: {err}')
         return None
 
 
+def process_batch(text_batch, model_type, embeddings_list, chunk_num, total_chunks):
+    print(f'[status] {timestamp_str()} - Processing chunk {chunk_num} of {total_chunks}')
+    result = send_request(text_batch, model_type)
+    if result:
+        if result[0]['status'] == 'error':
+            print(f'[error] {timestamp_str()} - Received error: {result[0]["message"]}')
+            return
+        else:
+            print(f'[status] {timestamp_str()} - Received embeddings: {len(result[0]["embeddings"])} ')
+            for text, em in zip(text_batch, result[0]['embeddings']):
+                metadata = {
+                    'status': result[0]['status'],
+                    'elapsed': result[0]['elapsed'],
+                    'model': result[0]['model']
+                }
+                embedding = Embedding(text=text, embedding=em, metadata=metadata)
+                embeddings_list.append(embedding.dict())
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
 
@@ -40,7 +74,7 @@ def send_request(text, model_type='local'):
     group.add_argument(
         '-f', '--file',
         type=argparse.FileType('r'),
-        help='text file to embed'
+        help='text file to embed (one text per line)'
     )
 
     parser.add_argument(
@@ -50,14 +84,65 @@ def send_request(text, model_type='local'):
         default='local'
     )
 
+    parser.add_argument(
+        '-o', '--output',
+        help='output file',
+        default='embeddings.json'
+    )
+
     args = parser.parse_args()
+    model_type = args.model
+    output_file = args.output
+    embeddings_list = []
+
+    if os.path.exists(output_file):
+        print(f'[error] {timestamp_str()} - Output file already exists')
+        sys.exit(1)
 
     if args.file:
-        text = args.file.read()
+        if not os.path.exists(args.file.name):
+            print(f'[error] {timestamp_str()} - File does not exist')
+            sys.exit(1)
+
+        print(f'[status] {timestamp_str()} - Processing file: {args.file.name}')
+
+        text_batch = []
+        chunk_size = 100
+        total_lines = sum(1 for _ in args.file)
+        args.file.seek(0)
+        total_chunks = (total_lines + chunk_size - 1) // chunk_size
+
+        print(f'[info] {timestamp_str()} - Total chunks: {total_chunks}')
+
+        chunk_num = 1
+
+        for line in args.file:
+            text = line.strip()
+            text_batch.append(text)
+            if len(text_batch) == chunk_size:
+                process_batch(text_batch, model_type, embeddings_list, chunk_num, total_chunks)
+                text_batch = []
+                chunk_num += 1
+
+        if text_batch:
+            process_batch(text_batch, model_type, embeddings_list, chunk_num, total_chunks)
     else:
+        print(f'[status] {timestamp_str()} - Processing text input')
         text = args.text
+        result = send_request([text], model_type)
+        if result:
+            for res in result:
+                metadata = {'status': res['status'], 'elapsed': res['elapsed'], 'model': res['model']}
+                embedding = Embedding(text=text, embedding=res['embedding'], metadata=metadata)
+                embeddings_list.append(embedding.dict())
+
+    try:
+        print(f'[status] {timestamp_str()} - Saving embeddings to {output_file}')
+        with open(output_file, 'w') as f:
+            json.dump(embeddings_list, f)
+
+        print(f'[status] {timestamp_str()} - Embeddings saved to embeddings.json')
+    except Exception as err:
+        print(f'[error] {timestamp_str()} - exception saving embeddings: {err}')
+        sys.exit(1)
 
-    model_type = args.model
-    result = send_request(text, model_type)
-    if result is not None:
-        print(result)