fix(kreuzberg): rewrite tool code to fix parsing errors

- Remove complex docstrings that caused parsing issues
- Use simple string concatenation instead of f-strings with special chars
- Remove duplicate code
- Keep functionality intact: single file, multi-file, list, search
This commit is contained in:
Folzo
2026-04-11 00:47:21 +02:00
parent 98b695bc31
commit e9d6de3513
+48 -122
View File
@@ -11,17 +11,11 @@
## Tool Code
```python
"""
Kreuzberg OCR Tool for OpenWebUI
Extracts text from scanned documents, images, and PDFs using Kreuzberg OCR.
Supports: PDF, JPEG, PNG, TIFF, DOCX, PPTX, XLSX, EPUB, etc.
"""
import os
import glob
import requests
import hashlib
import io
import json
from PIL import Image, ImageEnhance
@@ -31,8 +25,7 @@ class Tools:
self.uploads_dir = "/app/backend/data/uploads"
self.host_uploads_dir = "/home/effi/effi/data/uploads"
def _preprocess_image(self, image: Image.Image) -> Image.Image:
"""Preprocess image for better OCR quality (grayscale + contrast)"""
def _preprocess_image(self, image):
if image.mode in ('RGBA', 'LA', 'P'):
image = image.convert('RGBA').convert('RGB')
elif image.mode != 'RGB':
@@ -42,47 +35,24 @@ class Tools:
image = enhancer.enhance(2.0)
return image
def extract_file(self, filename: str) -> str:
"""
Extract text from a single file using Kreuzberg OCR.
Args:
filename: The filename shown in the chat (e.g., "document.pdf", "image.jpg")
or the full name like "UUID_name.pdf"
Returns:
Extracted text content with markdown formatting
"""
def extract_file(self, filename):
return self._extract_single_file(filename)
def extract_files(self, filenames: str) -> str:
"""
Extract text from multiple files, one at a time.
Args:
filenames: Comma-separated list of filenames (e.g., "doc1.pdf,doc2.jpg")
or JSON array string (e.g., '["doc1.pdf","doc2.jpg"]')
Returns:
Combined extracted text from all files, processed sequentially
"""
import json
def extract_files(self, filenames):
filenames = filenames.strip()
if not filenames:
return "❌ No filenames provided. Usage: extract_files(\"file1.pdf,file2.jpg\")"
return "[ERROR] No filenames provided. Usage: extract_files(\"file1.pdf,file2.jpg\")"
if filenames.startswith('['):
try:
file_list = json.loads(filenames)
except json.JSONDecodeError:
return f"❌ Invalid JSON array: {filenames}"
return "[ERROR] Invalid JSON array: " + filenames
else:
file_list = [f.strip() for f in filenames.split(',')]
if not file_list:
return "❌ No files to process"
return "[ERROR] No files to process"
if len(file_list) == 1:
return self._extract_single_file(file_list[0])
@@ -94,22 +64,19 @@ class Tools:
filename = filename.strip()
if not filename:
continue
if total > 1:
results.append(f"--- Document {i}/{total}: `{filename}` ---")
results.append("--- Document " + str(i) + "/" + str(total) + ": " + filename + " ---")
result = self._extract_single_file(filename)
results.append(result)
return "\n\n".join(results)
def _extract_single_file(self, filename: str) -> str:
"""Internal method to extract text from a single file"""
def _extract_single_file(self, filename):
try:
file_path = self._find_file(filename)
if not file_path:
return f"❌ File not found: {filename}\n\nTip: Use `list_recent_files()` to see available files."
return "[ERROR] File not found: " + filename + ". Use list_recent_files() to see available files."
mime_type = self._get_mime_type(file_path)
ext = os.path.splitext(file_path)[1].lower()
@@ -137,7 +104,7 @@ class Tools:
}
response = requests.post(
f"{self.kreuzberg_url}/v1/convert/file",
self.kreuzberg_url + "/v1/convert/file",
files=files,
data=data,
timeout=300
@@ -152,31 +119,30 @@ class Tools:
content = self._clean_content(content)
if len(content) > 15000:
content = content[:15000] + "\n\n... (document truncated - too long for context)"
content = content[:15000] + "\n\n... (document truncated)"
return f"📄 **{os.path.basename(file_path)}**\n\n{content}"
return "FILE: " + os.path.basename(file_path) + "\n\n" + content
else:
return f"⚠️ No text extracted from {os.path.basename(file_path)}\n\nThe file may be empty, protected, or contain only images without selectable text."
return "[WARNING] No text extracted from " + os.path.basename(file_path) + ". File may be empty or contain only images."
return f"⚠️ Unexpected response format from Kreuzberg: {str(result)[:500]}"
return "[WARNING] Unexpected response from Kreuzberg"
elif response.status_code == 404:
return f"❌ Kreuzberg endpoint not found. Is the server running at {self.kreuzberg_url}?"
return "[ERROR] Kreuzberg endpoint not found at " + self.kreuzberg_url
else:
return f"❌ Kreuzberg error ({response.status_code}): {response.text[:500]}"
return "[ERROR] Kreuzberg error (" + str(response.status_code) + ")"
except requests.Timeout:
return f"⏱️ Timeout - File is too large or Kreuzberg is overloaded. Try a smaller file."
return "[ERROR] Timeout - file is too large or Kreuzberg is overloaded"
except requests.ConnectionError:
return f"❌ Cannot connect to Kreuzberg at {self.kreuzberg_url}. Is the server running?"
return "[ERROR] Cannot connect to Kreuzberg at " + self.kreuzberg_url
except Exception as e:
return f"❌ Error: {str(e)}"
return "[ERROR] " + str(e)
def list_recent_files(self, count: int = 10) -> str:
"""List recent uploaded files"""
def list_recent_files(self, count=10):
try:
files = glob.glob(os.path.join(self.uploads_dir, "*"))
@@ -187,81 +153,70 @@ class Tools:
files = sorted(files, key=os.path.getmtime, reverse=True)[:count]
if not files:
return "📂 No files found in uploads directory."
return "[INFO] No files found in uploads directory."
result = f"📂 **{len(files)} recent files:**\n\n"
result = "[INFO] " + str(len(files)) + " recent files:\n\n"
for f in files:
name = os.path.basename(f)
size = os.path.getsize(f)
size_str = self._format_size(size)
result += f"- `{name}` ({size_str})\n"
result += "- " + name + " (" + size_str + ")\n"
result += "\n💡 Use `extract_file(\"filename\")` to OCR a file."
result += "\nUse extract_file(\"filename\") to OCR a file."
return result
except Exception as e:
return f"❌ Error listing files: {str(e)}"
return "[ERROR] " + str(e)
def search_files(self, pattern: str) -> str:
"""Search for files matching a pattern"""
def search_files(self, pattern):
try:
# Search in uploads
files = glob.glob(os.path.join(self.uploads_dir, f"*{pattern}*"))
files = glob.glob(os.path.join(self.uploads_dir, "*" + pattern + "*"))
if not files and os.path.exists(self.host_uploads_dir):
files = glob.glob(os.path.join(self.host_uploads_dir, f"*{pattern}*"))
files = glob.glob(os.path.join(self.host_uploads_dir, "*" + pattern + "*"))
files = [f for f in files if os.path.isfile(f)]
if not files:
return f"🔍 No files found matching: {pattern}"
return "[INFO] No files found matching: " + pattern
result = f"🔍 **{len(files)} files found:**\n\n"
for f in files[:20]: # Limit to 20
result = "[INFO] " + str(len(files)) + " files found:\n\n"
for f in files[:20]:
name = os.path.basename(f)
size = self._format_size(os.path.getsize(f))
result += f"- `{name}` ({size})\n"
result += "- " + name + " (" + size + ")\n"
return result
except Exception as e:
return f"❌ Error searching files: {str(e)}"
return "[ERROR] " + str(e)
def _find_file(self, filename: str) -> str:
"""Find file in uploads directory"""
# Clean up filename
def _find_file(self, filename):
filename = filename.strip().strip('"').strip("'")
# Try exact path
if os.path.exists(filename):
return filename
# Try in uploads directory with exact name
path = os.path.join(self.uploads_dir, filename)
if os.path.exists(path):
return path
# Try host mount path
path = os.path.join(self.host_uploads_dir, filename)
if os.path.exists(path):
return path
# Try glob patterns
for directory in [self.uploads_dir, self.host_uploads_dir]:
if not os.path.exists(directory):
continue
# UUID prefix match
pattern = os.path.join(directory, f"{filename}*")
pattern = os.path.join(directory, filename + "*")
matches = glob.glob(pattern)
if matches:
return matches[0]
# Filename contains match
pattern = os.path.join(directory, f"*{filename}*")
pattern = os.path.join(directory, "*" + filename + "*")
matches = glob.glob(pattern)
if matches:
# Prefer exact matches
for m in matches:
if filename in os.path.basename(m):
return m
@@ -269,8 +224,7 @@ class Tools:
return None
def _get_mime_type(self, file_path: str) -> str:
"""Determine MIME type from file extension"""
def _get_mime_type(self, file_path):
ext = os.path.splitext(file_path)[1].lower()
mime_types = {
'.pdf': 'application/pdf',
@@ -295,12 +249,10 @@ class Tools:
}
return mime_types.get(ext, 'application/octet-stream')
def _clean_content(self, content: str) -> str:
"""Clean up extracted content"""
def _clean_content(self, content):
if not content:
return ""
# Remove excessive whitespace
lines = content.split('\n')
cleaned_lines = []
prev_empty = False
@@ -315,64 +267,38 @@ class Tools:
cleaned_lines.append(stripped)
prev_empty = False
# Join and limit consecutive empty lines
result = '\n'.join(cleaned_lines)
# Remove very long lines (>200 chars wrap)
if len(result) > 500:
# Check if content looks garbled (lots of short fragments)
short_words = sum(1 for w in result.split() if len(w) <= 3)
total_words = len(result.split())
if total_words > 50 and short_words / total_words > 0.3:
result += "\n\n⚠️ Warning: This text may contain OCR errors due to complex formatting or handwriting."
result += "\n\n[WARNING] This text may contain OCR errors."
return result
def _format_size(self, size: int) -> str:
"""Format file size"""
def _format_size(self, size):
if size < 1024:
return f"{size} B"
return str(size) + " B"
elif size < 1024 * 1024:
return f"{size / 1024:.1f} KB"
return str(size / 1024) + " KB"
else:
return f"{size / (1024 * 1024):.1f} MB"
return str(size / (1024 * 1024)) + " MB"
# OpenWebUI Tool Functions
def extract_file(filename: str) -> str:
"""
Extract text from a single file using Kreuzberg OCR.
Args:
filename: The filename or UUID from the chat (e.g., "document.pdf", "abc123_document.pdf")
Returns:
Extracted text with markdown formatting
"""
def extract_file(filename):
tools = Tools()
return tools.extract_file(filename)
def extract_files(filenames: str) -> str:
"""
Extract text from multiple files, one at a time.
Args:
filenames: Comma-separated list (e.g., "doc1.pdf,doc2.jpg")
or JSON array (e.g., '["doc1.pdf","doc2.jpg"]')
Returns:
Combined extracted text from all files
"""
def extract_files(filenames):
tools = Tools()
return tools.extract_files(filenames)
def list_recent_files(count: int = 10) -> str:
"""List recent uploaded files"""
def list_recent_files(count=10):
tools = Tools()
return tools.list_recent_files(count)
def search_files(pattern: str) -> str:
"""Search for files by name pattern"""
def search_files(pattern):
tools = Tools()
return tools.search_files(pattern)
```