initial open-source release

llm_server/test_rag.py

@@ -0,0 +1,82 @@
+import requests
+import json
+
+# 服务地址
+BASE_URL = "http://0.0.0.0:8903"
+
+def print_separator(title):
+    print("\n" + "=" * 80)
+    print(f"  {title}")
+    print("=" * 80)
+
+def test_health():
+    print_separator("Health Check")
+    resp = requests.get(f"{BASE_URL}/health")
+    print(f"Status: {resp.status_code}")
+    print(f"Response: {resp.json()}")
+
+def test_embed():
+    print_separator("Embedding Test")
+    payload = {
+        "task_description": "Given a web search query, retrieve relevant passages that answer the query",
+        "queries": [
+            "What is the capital of France?",
+            "How to train a cat?"
+        ],
+        "documents": [
+            "Paris is the capital of France.",
+            "Cats can be trained using positive reinforcement."
+        ]
+    }
+    print(">>> Request Body:")
+    print(json.dumps(payload, indent=2, ensure_ascii=False))
+
+    resp = requests.post(f"{BASE_URL}/embed", json=payload)
+    print(f"<<< Status: {resp.status_code}")
+    if resp.status_code == 200:
+        data = resp.json()
+        # 只打印向量长度和前面几个值，避免刷屏
+        print(f"<<< Response Keys: {list(data.keys())}")
+        print(f"Number of query embeddings: {len(data['query_embeddings'])}")
+        print(f"Dimension of each query embedding: {len(data['query_embeddings'][0])}")
+        print(f"First query embedding (first 5 values): {data['query_embeddings'][0][:5]}")
+        print(f"Number of doc embeddings: {len(data['doc_embeddings'])}")
+        if data['doc_embeddings']:
+            print(f"Dimension of each doc embedding: {len(data['doc_embeddings'][0])}")
+            print(f"First doc embedding (first 5 values): {data['doc_embeddings'][0][:5]}")
+        else:
+            print("No document embeddings returned.")
+        # 如果需要完整输出，取消下面注释（会打印大量浮点数）
+        # print(f"<<< Full Response:\n{json.dumps(data, indent=2)}")
+    else:
+        print(f"Error: {resp.text}")
+
+def test_rerank():
+    print_separator("Reranker Test")
+    payload = {
+        "task_description": "Given a web search query, retrieve relevant passages that answer the query",
+        "query": "What is the capital of France?",
+        "documents": [
+            "Paris is the capital of France.",
+            "Cats are often difficult to train.",
+            "Lyon is a city in France."
+        ]
+    }
+    print(">>> Request Body:")
+    print(json.dumps(payload, indent=2, ensure_ascii=False))
+
+    resp = requests.post(f"{BASE_URL}/rerank", json=payload)
+    print(f"<<< Status: {resp.status_code}")
+    if resp.status_code == 200:
+        data = resp.json()
+        print(f"<<< Response Keys: {list(data.keys())}")
+        print(f"Scores: {data['scores']}")
+        # 打印完整响应
+        # print(f"<<< Full Response:\n{json.dumps(data, indent=2)}")
+    else:
+        print(f"Error: {resp.text}")
+
+if __name__ == "__main__":
+    test_health()
+    test_embed()
+    test_rerank()