Merge pull request #109 from GSA-TTS/track-costs

Track costs

backend/apps/socket/main.py

@@ -6,17 +6,6 @@
 from utils.utils import decode_token
 
 sio = socketio.AsyncServer(cors_allowed_origins=[], async_mode="asgi")
-
-# ENABLE_WEBSOCKET_SUPPORT = False
-
-# sio = socketio.AsyncServer(
-#     cors_allowed_origins=[],
-#     async_mode="asgi",
-#     transports=(["polling", "websocket"] if ENABLE_WEBSOCKET_SUPPORT else ["polling"]),
-#     allow_upgrades=ENABLE_WEBSOCKET_SUPPORT,
-#     always_connect=True,
-# )
-
 app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
 
 # Dictionary to maintain the user pool

backend/apps/webui/routers/documents.py

@@ -166,7 +166,7 @@ async def delete_doc_by_name(name: str, user=Depends(get_current_user)):
     log.info(f"Deleting document {name}")
     result = Documents.delete_doc_by_name(name)
 
-    # NOTE: storing the vectors less than creating them, let's skip while in beta
+    # NOTE: storing the vectors costs less than creating them, let's skip while in beta
     # success = VECTOR_CLIENT.delete(collection_name)
 
     return result

backend/config.py

@@ -5,12 +5,8 @@
 import pkgutil
 import chromadb
 
-# from chromadb import Settings
-# from redis.asyncio import Redis
-
-# from base64 import b64encode
 from bs4 import BeautifulSoup
-from typing import TypeVar, Generic  # , Union
+from typing import TypeVar, Generic
 from pydantic import BaseModel
 from typing import Optional
 
@@ -794,21 +790,8 @@ class BannerModel(BaseModel):
     os.environ.get("RAG_EMBEDDING_ENGINE", "openai"),
 )
 
-# Redis connection settings
-# REDIS_HOST = os.environ.get("REDIS_HOST", "localhost")
-# REDIS_PORT = int(os.environ.get("REDIS_PORT", "6379"))
-# REDIS_PASSWORD = os.environ.get("REDIS_PASSWORD", "")
-# REDIS_URL = f"redis://localhost:{REDIS_PORT}"
-# REDIS_ENABLED = True
-# REDIS_SSL = False
-# REDIS_SSL_CERT_REQS = None
-# REDIS_CONFIG = dict(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD)
-# REDIS_DB = int(os.environ.get("REDIS_DB", "0"))
-# REDIS_INDEX_NAME = os.environ.get("REDIS_INDEX_NAME", "document-index")
-# REDIS_PREFIX = os.environ.get("REDIS_PREFIX", "doc")
-# REDIS_VECTOR_DIM = 1536 if RAG_EMBEDDING_ENGINE == "openai" else 384
-
 VECTOR_STORE = os.environ.get("VECTOR_STORE", "postgres")
+# TODO: rip out chroma
 CHROMA_DATA_PATH = f"{DATA_DIR}/vector_db"
 CHROMA_TENANT = os.environ.get("CHROMA_TENANT", chromadb.DEFAULT_TENANT)
 CHROMA_DATABASE = os.environ.get("CHROMA_DATABASE", chromadb.DEFAULT_DATABASE)

backend/main.py

@@ -106,7 +106,6 @@
     VECTOR_CLIENT,
     REDIS_VL_SCHEMA,
     AppConfig,
-    # initialize_vector_client,
 )
 from constants import ERROR_MESSAGES
 
@@ -373,12 +372,7 @@ async def dispatch(self, request: Request, call_next):
             # If docs field is present, generate RAG completions
             if "docs" in data:
                 data = {**data}
-                log.error(
-                    f"docs: {data}\n"
-                    # + f"len messages: {len(data['messages'])}\n"
-                    # + f"{rag_app.state.EMBEDDING_FUNCTION}\n"
-                    # + f"{rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH}"
-                )
+                log.error(f"docs: {data}\n")
                 try:
                     log.error("About to call get_rag_context...")
                     rag_context, citations = await get_rag_context(

backend/requirements.txt

@@ -14,7 +14,9 @@ requests==2.32.2
 aiohttp==3.9.5
 peewee==3.17.5
 peewee-migrate==1.12.2
+psycopg==3.2.3
 psycopg2-binary==2.9.9
+pgvector==0.2.5
 PyMySQL==1.1.1
 bcrypt==4.1.3
 

src/lib/apis/rag/index.ts

@@ -168,7 +168,6 @@ export const uploadDocToVectorDB = async (token: string, collection_name: string
 	const data = new FormData();
 	data.append('file', file);
 	data.append('collection_name', collection_name);
-	console.log('collection_name', collection_name);
 
 	let error = null;