From 43500ce66f34d12948d0550ed37faa6ae1fa52c8 Mon Sep 17 00:00:00 2001 From: Christian Krakau-Louis Date: Fri, 28 Mar 2025 13:50:10 +0100 Subject: [PATCH] feat: integrate OpenAI API with configurable settings for model and base URL --- .env.demo | 2 ++ app/config.py | 2 ++ app/tasks/extract_metadata_with_gpt.py | 13 ++++++++----- app/tasks/refine_text_with_gpt.py | 16 ++++++++-------- 4 files changed, 20 insertions(+), 13 deletions(-) diff --git a/.env.demo b/.env.demo index 144b4ee4..ff890ae6 100644 --- a/.env.demo +++ b/.env.demo @@ -15,6 +15,8 @@ PAPERLESS_HOST=https://paperless.example.com AWS_ACCESS_KEY_ID="" AWS_SECRET_ACCESS_KEY="" OPENAI_API_KEY="" +OPENAI_BASE_URL=https://api.openai.com/v1 +OPENAI_MODEL=gpt-4o-mini PAPERLESS_NGX_API_TOKEN= DROPBOX_APP_KEY= DROPBOX_APP_SECRET= diff --git a/app/config.py b/app/config.py index 64774626..616203e4 100644 --- a/app/config.py +++ b/app/config.py @@ -13,6 +13,8 @@ class Settings(BaseSettings): redis_url: str s3_bucket_name: str openai_api_key: str + openai_base_url: str = "https://api.openai.com/v1" # Default to OpenAI's endpoint + openai_model: str = "gpt-4o-mini" # Default model workdir: str dropbox_app_key: str dropbox_app_secret: str diff --git a/app/tasks/extract_metadata_with_gpt.py b/app/tasks/extract_metadata_with_gpt.py index 5adf01c0..18c44aee 100644 --- a/app/tasks/extract_metadata_with_gpt.py +++ b/app/tasks/extract_metadata_with_gpt.py @@ -2,15 +2,19 @@ import json import re -from openai import OpenAI from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf # Import the shared Celery instance from app.celery_app import celery +import openai -client = OpenAI(api_key=settings.openai_api_key) +# Initialize OpenAI client dynamically +client = openai.OpenAI( + api_key=settings.openai_api_key, + base_url=settings.openai_base_url +) def extract_json_from_text(text): """ @@ -31,8 +35,7 @@ def extract_json_from_text(text): @celery.task(base=BaseTaskWithRetry) def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str): - """Uses OpenAI GPT-4o-mini to classify document metadata.""" - + """Uses OpenAI to classify document metadata.""" prompt = f""" You are a specialized document analyzer trained to extract structured metadata from documents. Your task is to analyze the given text and return a well-structured JSON object. @@ -68,7 +71,7 @@ Return only valid JSON with no additional commentary. try: print(f"[DEBUG] Sending classification request for {s3_filename}...") completion = client.chat.completions.create( - model="gpt-4o-mini", + model=settings.openai_model, messages=[ {"role": "system", "content": "You are an intelligent document classifier."}, {"role": "user", "content": prompt} diff --git a/app/tasks/refine_text_with_gpt.py b/app/tasks/refine_text_with_gpt.py index fe687b6b..950117c8 100644 --- a/app/tasks/refine_text_with_gpt.py +++ b/app/tasks/refine_text_with_gpt.py @@ -1,23 +1,23 @@ #!/usr/bin/env python3 from app.config import settings -from openai import OpenAI +import openai from app.tasks.retry_config import BaseTaskWithRetry # Import the shared Celery instance from app.celery_app import celery - - -client = OpenAI(api_key=settings.openai_api_key) +# Initialize OpenAI client dynamically +client = openai.OpenAI( + api_key=settings.openai_api_key, + base_url=settings.openai_base_url +) @celery.task(base=BaseTaskWithRetry) def refine_text_with_gpt(s3_filename: str, raw_text: str): - """Uses GPT to clean and refine OCR text.""" - - # Use the Chat Completions endpoint with 'messages' + """Uses OpenAI to clean and refine OCR text.""" response = client.chat.completions.create( - model="gpt-4", + model=settings.openai_model, messages=[ {"role": "system", "content": "Clean and format the following text. The idea is that the text you see comes from an OCR system and your task is to eliminate OCR errors. Keep the original language when doing so."}, {"role": "user", "content": raw_text}