Use JSON mode for custom llm

circlemind-ai · Nov 27, 2024 · e1f4322 · e1f4322
1 parent 3194f6e
commit e1f4322
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/examples/custom_llm.py b/examples/custom_llm.py
@@ -1,6 +1,8 @@
 """Example usage of GraphRAG with custom LLM and Embedding services compatible with the OpenAI API."""
+
 from typing import List
 
+import instructor
 from dotenv import load_dotenv
 
 from fast_graphrag import GraphRAG
@@ -19,7 +21,9 @@
     example_queries="\n".join(QUERIES),
     entity_types=ENTITY_TYPES,
     config=GraphRAG.Config(
-        llm_service=OpenAILLMService(model="your-llm-model", base_url="llm.api.url.com", api_key="your-api-key"),
+        llm_service=OpenAILLMService(
+            model="your-llm-model", base_url="llm.api.url.com", api_key="your-api-key", mode=instructor.Mode.JSON
+        ),
         embedding_service=OpenAIEmbeddingService(
             model="your-embedding-model",
             base_url="emb.api.url.com",