Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import requests | |
| import logging | |
| from typing import Optional | |
| from agents import function_tool | |
| from docx import Document | |
| import PyPDF2 | |
| # Set up logging | |
| logger = logging.getLogger(__name__) | |
| def read_document_data(query: str) -> str: | |
| """ | |
| Read and search for information from documents stored locally. | |
| Args: | |
| query: The search query or topic to look for in the documents | |
| Returns: | |
| The relevant content from the document(s) matching the query | |
| """ | |
| logger.info(f"TOOL CALL: read_document_data called with query='{query}'") | |
| local_content = _read_local_documents(query) | |
| if local_content: | |
| logger.info(f"TOOL RESULT: read_document_data found results") | |
| return local_content | |
| else: | |
| response = f"No relevant information found for query: '{query}'. Please check if documents are available." | |
| logger.info(f"TOOL RESULT: read_document_data found no results for query='{query}'") | |
| return response | |
| def _read_local_documents(query: str) -> Optional[str]: | |
| """Read from local PDF and DOCX files in the root directory.""" | |
| root_dir = os.path.dirname(os.path.dirname(__file__)) | |
| content_parts = [] | |
| # Try to read DOCX file | |
| docx_path = os.path.join(root_dir, "data.docx") | |
| if os.path.exists(docx_path): | |
| try: | |
| doc = Document(docx_path) | |
| full_text = [] | |
| for paragraph in doc.paragraphs: | |
| if paragraph.text.strip(): | |
| full_text.append(paragraph.text) | |
| docx_content = "\n".join(full_text) | |
| if docx_content: | |
| content_parts.append(f"[From data.docx]\n{docx_content}") | |
| except Exception as e: | |
| content_parts.append(f"Error reading data.docx: {str(e)}") | |
| # Try to read PDF files | |
| for file in os.listdir(root_dir): | |
| if file.endswith(".pdf"): | |
| pdf_path = os.path.join(root_dir, file) | |
| try: | |
| with open(pdf_path, "rb") as pdf_file: | |
| pdf_reader = PyPDF2.PdfReader(pdf_file) | |
| pdf_text = [] | |
| for page in pdf_reader.pages: | |
| text = page.extract_text() | |
| if text.strip(): | |
| pdf_text.append(text) | |
| if pdf_text: | |
| content_parts.append(f"[From {file}]\n" + "\n".join(pdf_text)) | |
| except Exception as e: | |
| content_parts.append(f"Error reading {file}: {str(e)}") | |
| return "\n\n".join(content_parts) if content_parts else None | |
| def _read_document_from_url(url: str, doc_name: str) -> Optional[str]: | |
| """Download and read a document (DOCX or PDF) from a URL.""" | |
| try: | |
| # Download the file from URL | |
| response = requests.get(url, timeout=30) | |
| response.raise_for_status() | |
| # Determine file type from URL | |
| if url.lower().endswith('.docx') or 'docx' in url.lower(): | |
| # Read DOCX from bytes | |
| doc = Document(io.BytesIO(response.content)) | |
| full_text = [] | |
| for paragraph in doc.paragraphs: | |
| if paragraph.text.strip(): | |
| full_text.append(paragraph.text) | |
| return "\n".join(full_text) | |
| elif url.lower().endswith('.pdf') or 'pdf' in url.lower(): | |
| # Read PDF from bytes | |
| pdf_reader = PyPDF2.PdfReader(io.BytesIO(response.content)) | |
| pdf_text = [] | |
| for page in pdf_reader.pages: | |
| text = page.extract_text() | |
| if text.strip(): | |
| pdf_text.append(text) | |
| return "\n".join(pdf_text) | |
| else: | |
| return f"Unsupported file type for URL: {url}" | |
| except Exception as e: | |
| raise Exception(f"Failed to download/read document from {url}: {str(e)}") | |
| def list_available_documents() -> str: | |
| """ | |
| List all available documents from local storage. | |
| Returns: | |
| A formatted list of available documents from local sources | |
| """ | |
| logger.info("TOOL CALL: list_available_documents called") | |
| # List local documents | |
| root_dir = os.path.dirname(os.path.dirname(__file__)) | |
| local_docs = [] | |
| if os.path.exists(os.path.join(root_dir, "data.docx")): | |
| local_docs.append("- data.docx") | |
| for file in os.listdir(root_dir): | |
| if file.endswith(".pdf"): | |
| local_docs.append(f"- {file}") | |
| if local_docs: | |
| response = "=== Local Documents ===\n" + "\n".join(local_docs) | |
| else: | |
| response = "No documents found in local storage." | |
| logger.info(f"TOOL RESULT: list_available_documents found {len(local_docs)} document(s)") | |
| return response | |