Skip to content

Commit

Permalink
Merge pull request #113 from GSA-TTS/track-costs
Browse files Browse the repository at this point in the history 
Track costs
  • Loading branch information
@jimmoffet
jimmoffet authored Nov 7, 2024
2 parents 13ec8e1 + 34876a4 commit 0c7712b
Show file tree
Hide file tree
Showing 4 changed files with 92 additions and 86 deletions.
22 changes: 11 additions & 11 deletions backend/apps/rag/clients/vector_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def has_collection(self, collection_name: str) -> bool:
return result and result[0]

def delete_collection(self, collection_name: str):
self.conn.execute(f"DROP TABLE IF EXISTS {collection_name}")
self.conn.execute(f'DROP TABLE IF EXISTS "{collection_name}"')

def search(
self, collection_name: str, vectors: list[list[float]], limit: int
Expand All @@ -59,7 +59,7 @@ def search(
# self.conn.execute("SET LOCAL hnsw.ef_search = 64")
query = f"""
SELECT id, content, (embedding <=> %s::vector) AS distance, metadata
FROM {collection_name}
FROM "{collection_name}"
ORDER BY distance
LIMIT %s;
"""
Expand All @@ -84,7 +84,7 @@ def query(
try:
where_clause = " AND ".join([f"{key} = %({key})s" for key in filter.keys()])
result = self.conn.execute(
f"SELECT id, content, metadata FROM {collection_name} WHERE {where_clause} LIMIT %(limit)s;",
f'SELECT id, content, metadata FROM "{collection_name}" WHERE {where_clause} LIMIT %(limit)s;',
{**filter, "limit": limit},
).fetchall()

Expand All @@ -102,7 +102,7 @@ def query(
def get(self, collection_name: str) -> Optional[GetResult]:
try:
result = self.conn.execute(
f"SELECT id, content, metadata FROM {collection_name};"
f'SELECT id, content, metadata FROM "{collection_name}";'
).fetchall()

if result:
Expand All @@ -120,13 +120,13 @@ def insert(self, collection_name: str, items: list[VectorItem]):
try:
if not self.has_collection(collection_name):
self.conn.execute(
f"CREATE TABLE {collection_name} "
f'CREATE TABLE "{collection_name}" '
f"(id bigserial PRIMARY KEY, content text, embedding vector(1536), metadata jsonb)"
)

for item in items:
self.conn.execute(
f"INSERT INTO {collection_name} (content, embedding, metadata) VALUES (%s, %s, %s)",
f'INSERT INTO "{collection_name}" (content, embedding, metadata) VALUES (%s, %s, %s)',
(item["text"], item["vector"], json.dumps(item["metadata"])),
)
except Exception as e:
Expand All @@ -136,14 +136,14 @@ def upsert(self, collection_name: str, items: list[VectorItem]):
try:
if not self.has_collection(collection_name):
self.conn.execute(
f"CREATE TABLE {collection_name} "
f'CREATE TABLE "{collection_name}" '
f"(id bigserial PRIMARY KEY, content text, embedding vector(1536), metadata jsonb)"
)

for item in items:
self.conn.execute(
f"""
INSERT INTO {collection_name} (content, embedding, metadata)
INSERT INTO "{collection_name}" (content, embedding, metadata)
VALUES (%s, %s, %s)
ON CONFLICT (id)
DO UPDATE SET content = EXCLUDED.content, embedding = EXCLUDED.embedding, metadata = EXCLUDED.metadata
Expand All @@ -164,18 +164,18 @@ def delete(
try:
if ids:
self.conn.execute(
f"DELETE FROM {collection_name} WHERE id IN %(ids)s;",
f'DELETE FROM "{collection_name}" WHERE id IN %(ids)s;',
{"ids": tuple(ids)},
)
elif filter:
where_clause = " AND ".join(
[f"{key} = %({key})s" for key in filter.keys()]
)
self.conn.execute(
f"DELETE FROM {collection_name} WHERE {where_clause};", filter
f'DELETE FROM "{collection_name}" WHERE {where_clause};', filter
)
elif not ids and not filter: # Check if both are empty
self.conn.execute(f"DROP TABLE IF EXISTS {collection_name};")
self.conn.execute(f'DROP TABLE IF EXISTS "{collection_name}";')
return True

except Exception as e:
Expand Down
112 changes: 59 additions & 53 deletions backend/apps/rag/main.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from asyncio import sleep
from fastapi import (
FastAPI,
Depends,
Expand Down Expand Up @@ -1211,61 +1212,66 @@ def store_text(

@app.get("/scan")
async def scan_docs_dir(user=Depends(get_admin_user)):
for path in Path(DOCS_DIR).rglob("./**/*"):
try:
if path.is_file() and not path.name.startswith("."):
tags = extract_folders_after_data_docs(path)
filename = path.name
file_content_type = mimetypes.guess_type(path)

f = open(path, "rb")
collection_name = calculate_sha256(f)[:63]
f.close()

loader, known_type = get_loader(
filename, file_content_type[0], str(path)
)
data = loader.load()

try:
result = await store_data_in_vector_db(data, collection_name)

if result:
sanitized_filename = sanitize_filename(filename)
doc = Documents.get_doc_by_name(sanitized_filename)

if doc is None:
doc = Documents.insert_new_doc(
user.id,
DocumentForm(
**{
"name": sanitized_filename,
"title": filename,
"collection_name": collection_name,
"filename": filename,
"content": (
json.dumps(
{
"tags": list(
map(
lambda name: {"name": name},
tags,
for dir in [DOCS_DIR, UPLOAD_DIR]:
for path in Path(dir).rglob("./**/*"):
try:
if path.is_file() and not path.name.startswith("."):
tags = extract_folders_after_data_docs(path)
filename = path.name
file_content_type = mimetypes.guess_type(path)

f = open(path, "rb")
log.info(f"Scanning {path}")
collection_name = calculate_sha256(f)[:63]
f.close()

loader, known_type = get_loader(
filename, file_content_type[0], str(path)
)
data = loader.load()

try:
await sleep(60)
result = await store_data_in_vector_db(data, collection_name)

if result:
sanitized_filename = sanitize_filename(filename)
doc = Documents.get_doc_by_name(sanitized_filename)

if doc is None:
doc = Documents.insert_new_doc(
user.id,
DocumentForm(
**{
"name": sanitized_filename,
"title": filename,
"collection_name": collection_name,
"filename": filename,
"content": (
json.dumps(
{
"tags": list(
map(
lambda name: {
"name": name
},
tags,
)
)
)
}
)
if len(tags)
else "{}"
),
}
),
)
except Exception as e:
log.exception(e)
pass
}
)
if len(tags)
else "{}"
),
}
),
)
except Exception as e:
log.exception(e)
pass

except Exception as e:
log.exception(e)
except Exception as e:
log.exception(e)

return True

Expand Down
4 changes: 2 additions & 2 deletions backend/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -399,14 +399,14 @@ async def dispatch(self, request: Request, call_next):

del data["docs"]

log.info(f"rag_context: {rag_context}, citations: {citations}")
log.debug(f"rag_context: {rag_context}, citations: {citations}")

if context != "":
system_prompt = rag_template(
rag_app.state.config.RAG_TEMPLATE, context, prompt
)

print(system_prompt)
log.info(system_prompt)

data["messages"] = add_or_update_system_message(
f"\n{system_prompt}", data["messages"]
Expand Down
40 changes: 20 additions & 20 deletions src/lib/components/workspace/Documents.svelte
View file
Original file line number Diff line number Diff line change
Expand Up @@ -463,26 +463,26 @@
</button>

<!-- <button
class="self-center w-fit text-sm px-2 py-2 border dark:border-gray-600 rounded-xl"
type="button"
on:click={() => {
console.log('download file');
}}
>
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 0 16 16"
fill="currentColor"
class="w-4 h-4"
>
<path
d="M8.75 2.75a.75.75 0 0 0-1.5 0v5.69L5.03 6.22a.75.75 0 0 0-1.06 1.06l3.5 3.5a.75.75 0 0 0 1.06 0l3.5-3.5a.75.75 0 0 0-1.06-1.06L8.75 8.44V2.75Z"
/>
<path
d="M3.5 9.75a.75.75 0 0 0-1.5 0v1.5A2.75 2.75 0 0 0 4.75 14h6.5A2.75 2.75 0 0 0 14 11.25v-1.5a.75.75 0 0 0-1.5 0v1.5c0 .69-.56 1.25-1.25 1.25h-6.5c-.69 0-1.25-.56-1.25-1.25v-1.5Z"
/>
</svg>
</button> -->
class="self-center w-fit text-sm px-2 py-2 border dark:border-gray-600 rounded-xl"
type="button"
on:click={() => {
console.log('download file');
}}
>
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 0 16 16"
fill="currentColor"
class="w-4 h-4"
>
<path
d="M8.75 2.75a.75.75 0 0 0-1.5 0v5.69L5.03 6.22a.75.75 0 0 0-1.06 1.06l3.5 3.5a.75.75 0 0 0 1.06 0l3.5-3.5a.75.75 0 0 0-1.06-1.06L8.75 8.44V2.75Z"
/>
<path
d="M3.5 9.75a.75.75 0 0 0-1.5 0v1.5A2.75 2.75 0 0 0 4.75 14h6.5A2.75 2.75 0 0 0 14 11.25v-1.5a.75.75 0 0 0-1.5 0v1.5c0 .69-.56 1.25-1.25 1.25h-6.5c-.69 0-1.25-.56-1.25-1.25v-1.5Z"
/>
</svg>
</button> -->

<button
class="self-center w-fit text-sm px-2 py-2 dark:text-gray-300 dark:hover:text-white hover:bg-black/5 dark:hover:bg-white/5 rounded-xl"
Expand Down

0 comments on commit 0c7712b

Please sign in to comment.