yangdx commited on
Commit
07c1c0f
·
1 Parent(s): 67da933

fix: add keyword_extraction param support for LLM func of API Server

Browse files
Files changed (1) hide show
  1. lightrag/api/lightrag_server.py +7 -0
lightrag/api/lightrag_server.py CHANGED
@@ -17,6 +17,7 @@ import argparse
17
  from typing import List, Any, Optional, Union, Dict
18
  from pydantic import BaseModel
19
  from lightrag import LightRAG, QueryParam
 
20
  from lightrag.api import __api_version__
21
  from lightrag.utils import EmbeddingFunc
22
  from enum import Enum
@@ -756,6 +757,9 @@ def create_app(args):
756
  keyword_extraction=False,
757
  **kwargs,
758
  ) -> str:
 
 
 
759
  return await openai_complete_if_cache(
760
  args.llm_model,
761
  prompt,
@@ -773,6 +777,9 @@ def create_app(args):
773
  keyword_extraction=False,
774
  **kwargs,
775
  ) -> str:
 
 
 
776
  return await azure_openai_complete_if_cache(
777
  args.llm_model,
778
  prompt,
 
17
  from typing import List, Any, Optional, Union, Dict
18
  from pydantic import BaseModel
19
  from lightrag import LightRAG, QueryParam
20
+ from lightrag.types import GPTKeywordExtractionFormat
21
  from lightrag.api import __api_version__
22
  from lightrag.utils import EmbeddingFunc
23
  from enum import Enum
 
757
  keyword_extraction=False,
758
  **kwargs,
759
  ) -> str:
760
+ keyword_extraction = kwargs.pop("keyword_extraction", None)
761
+ if keyword_extraction:
762
+ kwargs["response_format"] = GPTKeywordExtractionFormat
763
  return await openai_complete_if_cache(
764
  args.llm_model,
765
  prompt,
 
777
  keyword_extraction=False,
778
  **kwargs,
779
  ) -> str:
780
+ keyword_extraction = kwargs.pop("keyword_extraction", None)
781
+ if keyword_extraction:
782
+ kwargs["response_format"] = GPTKeywordExtractionFormat
783
  return await azure_openai_complete_if_cache(
784
  args.llm_model,
785
  prompt,