Update app.py via AI Editor

app.py

@@ -13,11 +13,6 @@ import openai
 import base64
 import datetime
 from werkzeug.utils import secure_filename
-import chromadb
-from chromadb.config import Settings
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(threadName)s %(message)s")
 logger = logging.getLogger("AskTricare")
@@ -27,64 +22,9 @@ SESSION_DATA = {}
 SESSION_LOCKS = {}
 SESSION_DIR_BASE = os.path.join(tempfile.gettempdir(), "asktricare_sessions")
 os.makedirs(SESSION_DIR_BASE, exist_ok=True)
-VECTOR_DB_DIR = os.path.join(os.getcwd(), "vector_db")
-DOCS_DIR = os.path.join(os.getcwd(), "doc")
-os.makedirs(DOCS_DIR, exist_ok=True)
-os.makedirs(VECTOR_DB_DIR, exist_ok=True)
 
 openai.api_key = os.environ.get("OPENAI_API_KEY")
 
-chroma_client = chromadb.Client(Settings(
-    chroma_db_impl="duckdb+parquet",
-    persist_directory=VECTOR_DB_DIR,
-))
-embeddings = OpenAIEmbeddings(model="text-embedding-ada-002", openai_api_key=openai.api_key)
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-
-def ingest_docs():
-    logger.info("Starting document ingestion...")
-    file_paths = []
-    for root, _, files in os.walk(DOCS_DIR):
-        for f in files:
-            if f.lower().endswith(('.txt', '.pdf', '.md', '.docx')):
-                file_paths.append(os.path.join(root, f))
-    documents = []
-    metadatas = []
-    ids = []
-    for path in file_paths:
-        try:
-            with open(path, "r", encoding="utf-8", errors="ignore") as infile:
-                content = infile.read()
-            chunks = text_splitter.split_text(content)
-            for idx, chunk in enumerate(chunks):
-                documents.append(chunk)
-                metadatas.append({"source": path, "chunk": idx})
-                ids.append(f"{os.path.basename(path)}_{idx}")
-        except Exception as e:
-            logger.error(f"Error ingesting {path}: {e}")
-    if documents:
-        vectordb = Chroma(
-            collection_name="asktricare",
-            embedding_function=embeddings,
-            persist_directory=VECTOR_DB_DIR,
-            client_settings=Settings(chroma_db_impl="duckdb+parquet", persist_directory=VECTOR_DB_DIR),
-        )
-        vectordb.add_texts(documents, metadatas=metadatas, ids=ids)
-        vectordb.persist()
-        logger.info(f"Ingested {len(documents)} chunks from {len(file_paths)} files.")
-    else:
-        logger.info("No new documents to ingest.")
-
-if not os.listdir(VECTOR_DB_DIR):
-    ingest_docs()
-
-vectordb = Chroma(
-    collection_name="asktricare",
-    embedding_function=embeddings,
-    persist_directory=VECTOR_DB_DIR,
-    client_settings=Settings(chroma_db_impl="duckdb+parquet", persist_directory=VECTOR_DB_DIR),
-)
-
 def get_session_id():
     sid = flask_request.cookies.get("asktricare_session_id")
     if not sid:
@@ -106,6 +46,7 @@ def get_session_state(session_id):
         SESSION_DATA[session_id] = {
             "messages": [],
             "uploads": [],
+            "openai_file_ids": [],
             "created": datetime.datetime.utcnow().isoformat()
         }
     return SESSION_DATA[session_id]
@@ -233,6 +174,27 @@ app.layout = html.Div([
     ], style={"display": "flex"})
 ])
 
+def _upload_file_to_openai(file_path, purpose="assistants"):
+    try:
+        with open(file_path, 'rb') as f:
+            res = openai.File.create(
+                file=f,
+                purpose=purpose
+            )
+        logger.info(f"Uploaded file to OpenAI: {res.id}")
+        return res.id
+    except Exception as e:
+        logger.error(f"Failed to upload file to OpenAI: {e}")
+        return None
+
+def _get_openai_file_ids(session_state):
+    return session_state.get("openai_file_ids", [])
+
+def _is_supported_doc(filename):
+    ext = os.path.splitext(filename)[1].lower()
+    # OpenAI supports: txt, pdf, docx, md for assistants file search
+    return ext in [".txt", ".pdf", ".md", ".docx"]
+
 @app.callback(
     Output("session-id", "data"),
     Input("url", "href"),
@@ -243,7 +205,6 @@ def assign_session_id(_):
     d = get_session_dir(sid)
     load_session_state(sid)
     logger.info(f"Assigned session id: {sid}")
-    resp = dash.no_update
     return sid
 
 @app.callback(
@@ -284,6 +245,13 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
                 with open(fp, "wb") as f:
                     f.write(base64.b64decode(data))
                 uploads.append({"name": fname, "is_img": is_img, "path": fp})
+                # If document is supported, upload to OpenAI
+                if _is_supported_doc(fname):
+                    file_id = _upload_file_to_openai(fp)
+                    if file_id:
+                        if "openai_file_ids" not in state:
+                            state["openai_file_ids"] = []
+                        state["openai_file_ids"].append(file_id)
             state["uploads"].extend(uploads)
             save_session_state(session_id)
             logger.info(f"Session {session_id}: Uploaded files {[u['name'] for u in uploads]}")
@@ -292,27 +260,32 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
             loading = True
             state["messages"].append({"role": "user", "content": user_input})
             try:
-                docs = []
-                try:
-                    retr = vectordb.similarity_search(user_input, k=3)
-                    docs = [d.page_content for d in retr]
-                except Exception as e:
-                    logger.warning(f"Vector search failed: {e}")
-                context = "\n\n".join(docs)
+                # Use OpenAI's file search tool via ChatCompletion if files exist
+                file_ids = _get_openai_file_ids(state)
                 system_prompt = load_system_prompt()
                 messages = [
                     {"role": "system", "content": system_prompt},
                 ]
                 for m in state["messages"]:
                     messages.append({"role": m["role"], "content": m["content"]})
-                if context.strip():
-                    messages.append({"role": "system", "content": f"Relevant reference material:\n{context}"})
-                response = openai.ChatCompletion.create(
-                    model="gpt-3.5-turbo",
-                    messages=messages,
-                    max_tokens=700,
-                    temperature=0.2,
-                )
+                if file_ids:
+                    # Use 'tools' for file_search (RAG) if supported
+                    response = openai.ChatCompletion.create(
+                        model="gpt-3.5-turbo-1106",
+                        messages=messages,
+                        tools=[{"type": "file_search"}],
+                        tool_choice="file_search",
+                        file_ids=file_ids,
+                        max_tokens=700,
+                        temperature=0.2,
+                    )
+                else:
+                    response = openai.ChatCompletion.create(
+                        model="gpt-3.5-turbo",
+                        messages=messages,
+                        max_tokens=700,
+                        temperature=0.2,
+                    )
                 reply = response.choices[0].message.content
                 state["messages"].append({"role": "assistant", "content": reply})
                 logger.info(f"Session {session_id}: User: {user_input} | Assistant: {reply}")