Update app.py

app.py

@@ -81,68 +81,45 @@ def load_resources():
 load_resources()
 
 def predict(vector):
-    """处理768维向量输入并返回答案"""
-    start_time = time.time()
-    print(f"输入向量类型: {type(vector)}, 长度: {len(vector) if hasattr(vector, '__len__') else 'N/A'}")
-
     try:
-        # 处理不同输入格式
-        if isinstance(vector, list) and len(vector) > 0 and isinstance(vector[0], list):
-            # 已经是正确的二维数组格式
-            pass
-        elif isinstance(vector, list) and all(isinstance(x, float) for x in vector):
-            # 一维列表输入 (API调用)
-            #vector = [vector]
-            vector = np.array(vector).reshape(1, -1).astype('float32')
-        else:
-            error_msg = "错误：输入格式无效"
-            print(error_msg)
-            return {"error": error_msg}
-
-        if len(vector) != 1 or len(vector[0]) != 768:
-            error_msg = f"错误：需要1x768的二维数组，收到{len(vector)}x{len(vector[0]) if vector else 0}"
-            print(error_msg)
-            return {"error": error_msg}
-
-        # 添加实际处理逻辑
-        #vector_array = np.array(vector, dtype=np.float32)
-        # 转换格式
+        # 确保向量格式正确
         query_vector = np.array(vector).astype('float32').reshape(1, -1)
+        
+        # FAISS搜索
         D, I = index.search(query_vector, k=3)
-
+        
+        # 构建结果 - 使用安全的列名访问
         results = []
         for i in range(3):
             try:
-                result = metadata.iloc[I[0][i]]
-                confidence = 1/(1+D[0][i])
-                results.append({
-                    "source": result['source'],
-                    "confidence": round(float(confidence), 2)
-                })
+                idx = I[0][i]
+                result = {
+                    "source": metadata.iloc[idx]["source"],
+                    # 安全访问content字段
+                    "content": metadata.iloc[idx].get("content", ""),
+                    "confidence": float(1/(1+D[0][i]))
+                }
+                results.append(result)
             except Exception as e:
                 print(f"结果处理错误: {str(e)}")
-                results.append({"error": f"结果 {i+1}: 数据获取失败"})
-
-        print(f"处理完成 | 耗时: {time.time()-start_time:.2f}秒")
+                results.append({
+                    "error": f"结果 {i+1}: 数据获取失败",
+                    "confidence": 0.0
+                })
         
-        # 返回标准化的JSON格式
         return {
-            # "status": "success",
-            # "results": results
-            "source": [metadata.iloc[i]["source"] for i in I[0]],
-            "content": [metadata.iloc[i]["content"] for i in I[0]],
-            "confidence": [float(1/(1+d)) for d in D[0]]
+            "status": "success",
+            "results": results
         }
-        
     except Exception as e:
-        import traceback
-        error_msg = f"处理错误: {str(e)}\n{traceback.format_exc()}"
-        print(error_msg)
-        return {
-            "status": "error",
-            "message": "处理错误，请重试或联系管理员"
-        }
-
+        return JSONResponse(
+            status_code=500,
+            content={
+                "status": "error",
+                "message": f"服务器内部错误: {str(e)}"
+            }
+        )
+        
 # 创建FastAPI应用
 app = FastAPI()