update

src/blackbox/chroma_upsert.py

@@ -13,6 +13,9 @@ from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
 from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 import chromadb
+        
+import os
+import tempfile
 
 from ..utils import chroma_setting
 
@@ -50,7 +53,7 @@ class ChromaUpsert(Blackbox):
                 "collection_id": "123",
                 "action": "upsert",
                 "content": "file_name or string",
-                "answer": "success, collection has 100 documents.",
+                "answer": "collection 123 has 12472 documents. /tmp/Cheap and Quick：Efficient Vision-Language Instruction Tuning for Large Language Models.pdf ids is 0~111",
             },
         ]
 
@@ -59,10 +62,8 @@ class ChromaUpsert(Blackbox):
         else:
             collection_id = "123"
 
-        print("file: ",file)
-        print("file name: ",file.filename)
         if file is not None:
-            file_type = file.filename.split(".")[-1]
+            file_type = file.split(".")[-1]
             print("file_type: ",file_type)
             if file_type == "pdf":
                 loader = PyPDFLoader(file)
@@ -80,7 +81,6 @@ class ChromaUpsert(Blackbox):
                 loader = UnstructuredExcelLoader(file)
             
 
-            loader = PyPDFLoader(file)
             documents = loader.load()
             text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=0)
 
@@ -89,18 +89,28 @@ class ChromaUpsert(Blackbox):
             ids = [str(file)+str(i) for i in range(len(docs))]
 
             Chroma.from_documents(documents=docs, embedding=self.embedding_model, ids=ids, collection_name=collection_id, client=self.client)
+            
+            collection_number = self.client.get_collection(collection_id).count()
+            response_file = f"collection {collection_id} has {collection_number} documents. {file} ids is 0~{len(docs)-1}"
 
         if string is not None:
             # 生成一个新的id  ids_string: 1
-            ids =  setting.ChromaSetting.string_ids[0] + 1
+            # ids =  setting.ChromaSetting.string_ids[0] + 1
+            ids =  "1"
 
             Chroma.from_texts(texts=[string], embedding=self.embedding_model, ids=[ids], collection_name=collection_id, client=self.client)
 
 
-        collection_number = self.client.get_collection(collection_id).count()
-        response = f"collection {collection_id} has {collection_number} documents."
+            collection_number = self.client.get_collection(collection_id).count()
+            response_string = f"collection {collection_id} has {collection_number} documents. {string} ids is {ids}"
         
-        return response
+        
+        if file is not None and string is not None:
+            return response_file + " \n and " + response_string
+        elif file is not None and string is None:
+            return response_file
+        elif file is None and string is not None:
+            return response_string
 
 
 
@@ -117,18 +127,20 @@ class ChromaUpsert(Blackbox):
 
         if user_file is None and user_string is None:
             return JSONResponse(content={"error": "file or string is required"}, status_code=status.HTTP_400_BAD_REQUEST)
-        
-        # data = await user_file.read()
-        # with open(f'./{data.filename}', 'wb') as f:
-        #     f.write(content)
 
-        loader = PyPDFLoader(f'./{user_file.filename}')
-        documents = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=0)
-        docs = text_splitter.split_documents(documents)
-        print("docs: ",docs)
+        if user_file is not None:
+            pdf_bytes = await user_file.read()
+
+            custom_filename = user_file.filename
+            # 获取系统的临时目录路径
+            safe_filename = os.path.join(tempfile.gettempdir(), os.path.basename(custom_filename))
+        
+            with open(safe_filename, "wb") as f:
+                f.write(pdf_bytes)
+        else:
+            safe_filename = None
 
 
         return JSONResponse(
-            content={"response": self.processing(user_collection_id, user_file, user_string, user_context, user_setting)},
+            content={"response": self.processing(user_collection_id, safe_filename, user_string, user_context, user_setting)},
             status_code=status.HTTP_200_OK)