nw9663644-eng commented on code in PR #356: URL: https://github.com/apache/hugegraph-ai/pull/356#discussion_r3348486259
########## hugegraph-llm/src/tests/document/test_vector_index_utils.py: ########## @@ -0,0 +1,138 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import gradio as gr +import pytest +from docx import Document + +from hugegraph_llm.utils import vector_index_utils +from hugegraph_llm.utils.vector_index_utils import read_documents + + +def _build_pdf(content_stream: bytes) -> bytes: + objects = [ + b"<< /Type /Catalog /Pages 2 0 R >>", + b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>", + ( + b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] " + b"/Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>" + ), + b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", + (b"<< /Length " + str(len(content_stream)).encode() + b" >>\nstream\n" + content_stream + b"\nendstream"), + ] + + pdf = b"%PDF-1.4\n" + offsets = [] + for index, obj in enumerate(objects, start=1): + offsets.append(len(pdf)) + pdf += f"{index} 0 obj\n".encode() + pdf += obj + b"\nendobj\n" + + xref_offset = len(pdf) + pdf += f"xref\n0 {len(objects) + 1}\n".encode() + pdf += b"0000000000 65535 f \n" + for offset in offsets: + pdf += f"{offset:010d} 00000 n \n".encode() + + pdf += (f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\nstartxref\n{xref_offset}\n%%EOF\n").encode() + return pdf + + +def test_read_documents_reads_txt_file(tmp_path): + txt_path = tmp_path / "sample.txt" + txt_path.write_text("hello hugegraph", encoding="utf-8") + + result = read_documents([SimpleNamespace(name=str(txt_path))], "") + + assert result == ["hello hugegraph"] + + +def test_read_documents_reads_pdf_file(tmp_path): Review Comment: Thanks for the review. I added contract-level tests for both upload entrypoints: - build_vector_index() now verifies that a PDF upload is parsed and forwarded to the scheduler as extracted text. - extract_graph() now verifies that a PDF upload is parsed and forwarded to the graph extraction scheduler flow with the schema and prompt preserved. Local checks: - uv run pytest hugegraph-llm/src/tests/document/test_vector_index_utils.py - uv run ruff format --check . -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
