[add test][test llm and embedding api]

Author: lightislost

examples/ekg_examples/start.py

@@ -1,23 +1,30 @@
 import time, sys
-st = time.time()
 import os
 import yaml
 import requests 
 from typing import List
 from loguru import logger
-import tqdm
+from tqdm import tqdm
 from concurrent.futures import ThreadPoolExecutor
 
-print(time.time()-st)
 from langchain.llms.base import LLM
 from langchain.embeddings.base import Embeddings
-print(time.time()-st)
 src_dir = os.path.join(
     os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 )
-print(src_dir)
 sys.path.append(src_dir)
 
+try:
+    import os, sys
+    src_dir = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    )
+    sys.path.append(src_dir)
+    import test_config
+except Exception as e:
+    # set your config
+    logger.error(f"{e}")
+
 from muagent.schemas.db import *
 from muagent.llm_models.llm_config import EmbedConfig, LLMConfig
 from muagent.service.ekg_construct.ekg_construct_base import EKGConstructService
@@ -30,7 +37,7 @@ class CustomLLM(LLM, BaseModel):
     model_name: str = "qwen2:1b"
     model_type: str = "ollama"
     api_key: str = ""
-    stop: str = ""
+    stop: str = None
     temperature: float = 0.3
     top_k: int = 50
     top_p: float = 0.95
@@ -43,10 +50,8 @@ def params(self):
             if k in keys}
 
     def update_params(self, **kwargs):
-        logger.debug(f"{kwargs}")
         # 更新属性
         for key, value in kwargs.items():
-            logger.debug(f"{key}, {value}")
             setattr(self, key, value)
 
     def _llm_type(self, *args):
@@ -114,10 +119,8 @@ def params(self):
         }
 
     def update_params(self, **kwargs):
-        logger.debug(f"{kwargs}")
         # 更新属性
         for key, value in kwargs.items():
-            logger.debug(f"{key}, {value}")
             setattr(self, key, value)
 
     def _get_sentence_emb(self, sentence: str) -> dict:
@@ -133,6 +136,8 @@ def _get_sentence_emb(self, sentence: str) -> dict:
             return r.json()
         elif self.embedding_type == "openai":
             from muagent.llm_models.get_embedding import get_embedding
+            os.environ["OPENAI_API_KEY"] = self.api_key
+            os.environ["API_BASE_URL"] = self.url
             embed_config = EmbedConfig(
                 embed_engine="openai",
                 api_key=self.api_key,

muagent/httpapis/ekg_construct/api.py

@@ -2,6 +2,7 @@
 from typing import Dict
 import asyncio
 import uvicorn
+from loguru import logger
 
 from muagent.service.ekg_construct.ekg_construct_base import EKGConstructService
 from muagent.schemas.apis.ekg_api_schema import *
@@ -18,6 +19,24 @@ def init_app(llm, embeddings):
     async def llm_params():
         return llm.params()
 
+    # ~/llm/params
+    @app.post("/llm/generate", response_model=LLMResponse)
+    async def llm_predict(request: LLMRequest):
+        # 添加预测逻辑的代码
+        errorMessage = "ok"
+        successCode = True
+        try:
+            answer = llm.predict(request.text, request.stop)
+        except Exception as e:
+            errorMessage = str(e)
+            successCode = False
+            answer = "error"
+            
+        return LLMResponse(
+            successCode=successCode, errorMessage=errorMessage,
+            answer=answer
+        )
+
     # ~/llm/params/update
     @app.post("/llm/params/update", response_model=EKGResponse)
     async def update_llm_params(kwargs: Dict):
@@ -55,6 +74,23 @@ async def update_embedding_params(kwargs: Dict):
             successCode=successCode, errorMessage=errorMessage,
         )
 
+    @app.post("/embeddings/generate", response_model=EmbeddingsResponse)
+    async def embedding_predict(request: EmbeddingsRequest):
+        # 添加预测逻辑的代码
+        errorMessage = "ok"
+        successCode = True
+        try:
+            embeddings_list = embeddings.embed_documents(request.texts)
+        except Exception as e:
+            logger.exception(e)
+            errorMessage = str(e)
+            successCode = False
+            embeddings_list = []
+            
+        return EmbeddingsResponse(
+            successCode=successCode, errorMessage=errorMessage,
+            embeddings=embeddings_list
+        )
     # # ~/ekg/text2graph
     # @app.post("/ekg/text2graph", response_model=EKGGraphResponse)
     # async def text2graph(request: EKGT2GRequest):

muagent/schemas/apis/ekg_api_schema.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from typing import List, Dict
+from typing import List, Dict, Optional
 from enum import Enum
 
 from muagent.schemas.common import GNode, GEdge
@@ -12,6 +12,24 @@ class EKGResponse(BaseModel):
     errorMessage: str
 
 
+# embeddings
+class EmbeddingsResponse(EKGResponse):
+    successCode: int
+    errorMessage: str
+    embeddings: List[List[float]]
+
+class EmbeddingsRequest(BaseModel):
+    texts: List[str]
+
+class LLMRequest(BaseModel):
+    text: str
+    stop: Optional[str]
+
+class LLMResponse(EKGResponse):
+    successCode: int
+    errorMessage: str
+    answer: str
+
 # text2graph
 class EKGT2GRequest(BaseModel):
     text: str
@@ -81,7 +99,7 @@ class LLMParamsResponse(BaseModel):
     model_name: str
     model_type: str
     api_key: str
-    stop: str
+    stop: Optional[str] = None
     temperature: float
     top_k: int
     top_p: float