diff --git a/examples/lightrag_api_openai_compatible_demo.py b/examples/lightrag_api_openai_compatible_demo.py index 8173dc5b..f55b9ce1 100644 --- a/examples/lightrag_api_openai_compatible_demo.py +++ b/examples/lightrag_api_openai_compatible_demo.py @@ -24,6 +24,10 @@ EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-3-large") print(f"EMBEDDING_MODEL: {EMBEDDING_MODEL}") EMBEDDING_MAX_TOKEN_SIZE = int(os.environ.get("EMBEDDING_MAX_TOKEN_SIZE", 8192)) print(f"EMBEDDING_MAX_TOKEN_SIZE: {EMBEDDING_MAX_TOKEN_SIZE}") +BASE_URL = int(os.environ.get("BASE_URL", "https://api.openai.com/v1")) +print(f"BASE_URL: {BASE_URL}") +API_KEY = int(os.environ.get("API_KEY", "xxxxxxxx")) +print(f"API_KEY: {API_KEY}") if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) @@ -36,10 +40,12 @@ async def llm_model_func( prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs ) -> str: return await openai_complete_if_cache( - LLM_MODEL, - prompt, + model=LLM_MODEL, + prompt=prompt, system_prompt=system_prompt, history_messages=history_messages, + base_url=BASE_URL, + api_key=API_KEY, **kwargs, ) @@ -49,8 +55,10 @@ async def llm_model_func( async def embedding_func(texts: list[str]) -> np.ndarray: return await openai_embed( - texts, + texts=texts, model=EMBEDDING_MODEL, + base_url=BASE_URL, + api_key=API_KEY, ) diff --git a/lightrag/kg/jsondocstatus_impl.py b/lightrag/kg/jsondocstatus_impl.py index 675cf643..fe67b830 100644 --- a/lightrag/kg/jsondocstatus_impl.py +++ b/lightrag/kg/jsondocstatus_impl.py @@ -109,6 +109,22 @@ class JsonDocStatusStorage(DocStatusStorage): if v["status"] == DocStatus.PENDING } + async def get_processed_docs(self) -> dict[str, DocProcessingStatus]: + """Get all processed documents""" + return { + k: DocProcessingStatus(**v) + for k, v in self._data.items() + if v["status"] == DocStatus.PROCESSED + } + + async def get_processing_docs(self) -> dict[str, DocProcessingStatus]: + """Get all processing documents""" + return { + k: DocProcessingStatus(**v) + for k, v in self._data.items() + if v["status"] == DocStatus.PROCESSING + } + async def index_done_callback(self): """Save data to file after indexing""" write_json(self._data, self._file_name) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index eff4614d..5ec0dbeb 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -543,7 +543,7 @@ class LightRAG: new_docs = {doc_id: new_docs[doc_id] for doc_id in unique_new_doc_ids} if not new_docs: - logger.info("All documents have been processed or are duplicates") + logger.info("No new unique documents were found.") return # 4. Store status document @@ -560,15 +560,16 @@ class LightRAG: each chunk for entity and relation extraction, and updating the document status. - 1. Get all pending and failed documents + 1. Get all pending, failed, and abnormally terminated processing documents. 2. Split document content into chunks 3. Process each chunk for entity and relation extraction 4. Update the document status """ - # 1. get all pending and failed documents + # 1. Get all pending, failed, and abnormally terminated processing documents. to_process_docs: dict[str, DocProcessingStatus] = {} - # Fetch failed documents + processing_docs = await self.doc_status.get_processing_docs() + to_process_docs.update(processing_docs) failed_docs = await self.doc_status.get_failed_docs() to_process_docs.update(failed_docs) pendings_docs = await self.doc_status.get_pending_docs() @@ -599,6 +600,7 @@ class LightRAG: doc_status_id: { "status": DocStatus.PROCESSING, "updated_at": datetime.now().isoformat(), + "content": status_doc.content, "content_summary": status_doc.content_summary, "content_length": status_doc.content_length, "created_at": status_doc.created_at, @@ -635,6 +637,10 @@ class LightRAG: doc_status_id: { "status": DocStatus.PROCESSED, "chunks_count": len(chunks), + "content": status_doc.content, + "content_summary": status_doc.content_summary, + "content_length": status_doc.content_length, + "created_at": status_doc.created_at, "updated_at": datetime.now().isoformat(), } } @@ -648,6 +654,10 @@ class LightRAG: doc_status_id: { "status": DocStatus.FAILED, "error": str(e), + "content": status_doc.content, + "content_summary": status_doc.content_summary, + "content_length": status_doc.content_length, + "created_at": status_doc.created_at, "updated_at": datetime.now().isoformat(), } } diff --git a/lightrag/llm/openai.py b/lightrag/llm/openai.py index 535d665c..2a2ba0fd 100644 --- a/lightrag/llm/openai.py +++ b/lightrag/llm/openai.py @@ -103,17 +103,17 @@ async def openai_complete_if_cache( ) -> str: if history_messages is None: history_messages = [] - if api_key: - os.environ["OPENAI_API_KEY"] = api_key + if not api_key: + api_key = os.environ["OPENAI_API_KEY"] default_headers = { "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}", "Content-Type": "application/json", } openai_async_client = ( - AsyncOpenAI(default_headers=default_headers) + AsyncOpenAI(default_headers=default_headers, api_key=api_key) if base_url is None - else AsyncOpenAI(base_url=base_url, default_headers=default_headers) + else AsyncOpenAI(base_url=base_url, default_headers=default_headers, api_key=api_key) ) kwargs.pop("hashing_kv", None) kwargs.pop("keyword_extraction", None) @@ -294,17 +294,17 @@ async def openai_embed( base_url: str = None, api_key: str = None, ) -> np.ndarray: - if api_key: - os.environ["OPENAI_API_KEY"] = api_key + if not api_key: + api_key = os.environ["OPENAI_API_KEY"] default_headers = { "User-Agent": f"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) LightRAG/{__api_version__}", "Content-Type": "application/json", } openai_async_client = ( - AsyncOpenAI(default_headers=default_headers) + AsyncOpenAI(default_headers=default_headers, api_key=api_key) if base_url is None - else AsyncOpenAI(base_url=base_url, default_headers=default_headers) + else AsyncOpenAI(base_url=base_url, default_headers=default_headers, api_key=api_key) ) response = await openai_async_client.embeddings.create( model=model, input=texts, encoding_format="float"