85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
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 []
|
|
|