RAG Business in progress

lib/ragbusiness.py

@@ -0,0 +1,84 @@
+import requests, gc, lancedb, pyarrow, uuid
+from datetime import datetime
+
+def _schema_domainlogic(vector_size):
+    return pyarrow.schema([
+        pyarrow.field('timestamp',          pyarrow.timestamp('ms'),                       metadata={'description': 'When record created'}),
+        pyarrow.field('content',            pyarrow.string(),                              metadata={'description': 'Content'}),
+        pyarrow.field('category',           pyarrow.string(),                              metadata={'description': 'Category. Comma separated.'}),
+        pyarrow.field('vector_context',     pyarrow.list_(pyarrow.float32(), vector_size), metadata={'description': 'Vector data of combined data'}),
+    ])
+
+_TABLE_SCHEMAS = {
+    'knowledge_domainlogic':  _schema_domainlogic,
+}
+
+def db_init(db_path, vector_size):
+    db        = lancedb.connect(db_path)
+    existing  = db.table_names()
+    for name, schema_fn in _TABLE_SCHEMAS.items():
+        if name not in existing:
+            db.create_table(name, schema=schema_fn(vector_size))
+            print(f"[ragroleplay] Created table: {name}", flush=True)
+    del db
+    gc.collect()
+    
+def table_ensure(db_path, table_name, vector_size):
+    db        = lancedb.connect(db_path)
+    existing  = db.table_names()
+    if table_name not in existing:
+        schema_fn = _TABLE_SCHEMAS.get(table_name)
+        if schema_fn:
+            db.create_table(table_name, schema=schema_fn(vector_size))
+    del db
+    gc.collect()
+
+def _uuid(val):
+    return uuid.UUID(val) if isinstance(val, str) else val
+
+def embed_text(url, model, text):
+    response  = requests.post(url=url, json={"model": model, "input": text} )
+    data      = response.json()
+    return data["embeddings"][0]
+
+def domainlogic_store(db_path, model_url, model_name, payload):
+    try:
+        db      = lancedb.connect(db_path)
+        table   = db.open_table("knowledge_domainlogic")
+        
+        vector  = embed_text(model_url, model_name, f'Category: {payload["category"]}. \n{payload["content"]}')
+        
+        record  = {
+            "timestamp"       : datetime.now(),
+            "content"         : payload["content"   ],
+            "category"        : payload["category"  ],
+            "vector_context"  : vector
+        }
+        
+        table.add([record])
+        del table
+        del db
+        gc.collect()
+        
+        return "success"
+    
+    except Exception as e:
+        return e
+
+def domainlogic_load(db_path, payload):
+    try:
+        db       = lancedb.connect(db_path)
+        table    = db.open_table("knowledge_domainlogic")
+        
+        results  = table.search().to_list()
+        
+        del table
+        del db
+        gc.collect()
+        
+        return results
+    
+    except Exception as e:
+        print(f"error: {e}")
+        return []
+