Add support for third openai agent service

README.md

@@ -155,6 +155,12 @@ Create a `.env` file in the root directory and add your API key:
 CHATGPT_API_KEY=your_openai_key_here
 ```
 
+If you are using the OpenAI Agent Service, you can edit the `.env` file and add your Agent Service base address:
+
+```bash
+CHATGPT_API_BASE=your_openai_agent_service_base_address_here
+```
+
 ### 3. Run PageIndex on your PDF
 
 ```bash

pageindex/page_index.py

@@ -322,7 +322,7 @@ def toc_transformer(toc_content, model=None):
         if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)
 
 
-    last_complete = json.loads(last_complete)
+    last_complete = extract_json(last_complete)
 
     cleaned_response=convert_page_to_int(last_complete['table_of_contents'])
     return cleaned_response

pageindex/utils.py

@@ -18,6 +18,7 @@
 from types import SimpleNamespace as config
 
 CHATGPT_API_KEY = os.getenv("CHATGPT_API_KEY")
+CHATGPT_API_BASE = os.getenv("CHATGPT_API_BASE") or None
 
 def count_tokens(text, model=None):
     if not text:
@@ -28,7 +29,7 @@ def count_tokens(text, model=None):
 
 def ChatGPT_API_with_finish_reason(model, prompt, api_key=CHATGPT_API_KEY, chat_history=None):
     max_retries = 10
-    client = openai.OpenAI(api_key=api_key)
+    client = openai.OpenAI(api_key=api_key, base_url=CHATGPT_API_BASE)
     for i in range(max_retries):
         try:
             if chat_history:
@@ -60,7 +61,7 @@ def ChatGPT_API_with_finish_reason(model, prompt, api_key=CHATGPT_API_KEY, chat_
 
 def ChatGPT_API(model, prompt, api_key=CHATGPT_API_KEY, chat_history=None):
     max_retries = 10
-    client = openai.OpenAI(api_key=api_key)
+    client = openai.OpenAI(api_key=api_key, base_url=CHATGPT_API_BASE)
     for i in range(max_retries):
         try:
             if chat_history:
@@ -91,7 +92,7 @@ async def ChatGPT_API_async(model, prompt, api_key=CHATGPT_API_KEY):
     messages = [{"role": "user", "content": prompt}]
     for i in range(max_retries):
         try:
-            async with openai.AsyncOpenAI(api_key=api_key) as client:
+            async with openai.AsyncOpenAI(api_key=api_key, base_url=CHATGPT_API_BASE) as client:
                 response = await client.chat.completions.create(
                     model=model,
                     messages=messages,