From 3316534491478031aa2b58a2be39de824b3499e8 Mon Sep 17 00:00:00 2001 From: jayyuwriter Date: Mon, 10 Feb 2025 10:22:42 -0800 Subject: [PATCH] Added WriterQuestionToKG and tests block for querying knowledge graphs --- src/writer/blocks/__init__.py | 2 + src/writer/blocks/writerquestiontokg.py | 77 ++++++++++ .../backend/blocks/test_writerquestiontokg.py | 145 ++++++++++++++++++ 3 files changed, 224 insertions(+) create mode 100644 src/writer/blocks/writerquestiontokg.py create mode 100644 tests/backend/blocks/test_writerquestiontokg.py diff --git a/src/writer/blocks/__init__.py b/src/writer/blocks/__init__.py index 66bfdece4..02212f1b6 100644 --- a/src/writer/blocks/__init__.py +++ b/src/writer/blocks/__init__.py @@ -14,6 +14,7 @@ from writer.blocks.writercompletion import WriterCompletion from writer.blocks.writerinitchat import WriterInitChat from writer.blocks.writernocodeapp import WriterNoCodeApp +from writer.blocks.writerquestiontokg import WriterQuestionToKG SetState.register("workflows_setstate") WriterClassification.register("workflows_writerclassification") @@ -31,3 +32,4 @@ ReturnValue.register("workflows_returnvalue") WriterInitChat.register("workflows_writerinitchat") WriterAddToKG.register("workflows_writeraddtokg") +WriterQuestionToKG.register("workflows_writerquestiontokg") \ No newline at end of file diff --git a/src/writer/blocks/writerquestiontokg.py b/src/writer/blocks/writerquestiontokg.py new file mode 100644 index 000000000..5cc2b6e12 --- /dev/null +++ b/src/writer/blocks/writerquestiontokg.py @@ -0,0 +1,77 @@ +from writer.abstract import register_abstract_template +from writer.blocks.base_block import WorkflowBlock +from writer.ss_types import AbstractTemplate + + +class WriterQuestionToKG(WorkflowBlock): + @classmethod + def register(cls, type: str): + super(WriterQuestionToKG, cls).register(type) + register_abstract_template(type, AbstractTemplate( + baseType="workflows_node", + writer={ + "name": "Question Knowledge Graph", + "description": "Ask a question to the knowledge graph.", + "category": "Writer", + "fields": { + "graphId": { + "name": "Graph ids", + "type": "Text", + "desc": "The ids for existing knowledge graphs. For multiple graphs, provide comma-separated UUIDs (e.g., 123e4567-e89b-12d3-a456-426614174000, 550e8400-e29b-41d4-a716-446655440000)", + "validator": { + "type": "string", + }, + }, + "question": { + "name": "Question", + "type": "Text", + "desc": "The question to ask the knowledge graph.", + }, + "subqueries": { + "name": "Subqueries", + "type": "Text", + "desc": "Specify whether to include subqueries.", + "default": "no", + "options": { + "yes": "Yes", + "no": "No" + } + } + }, + "outs": { + "success": { + "name": "Success", + "description": "If the execution was successful.", + "style": "success", + }, + "error": { + "name": "Error", + "description": "If the function raises an Exception.", + "style": "error", + }, + }, + } + )) + + def run(self): + try: + import writer.ai + + graph_ids_str = self._get_field("graphId", required=True) + graph_ids = [id.strip() for id in graph_ids_str.split(',')] + question = self._get_field("question", required=True) + stream = self._get_field("stream", default_field_value="no") == "yes" + subqueries = self._get_field("subqueries", default_field_value="no") == "yes" + + client = writer.ai.WriterAIManager.acquire_client() + response = client.graphs.question( + graph_ids=graph_ids, + question=question, + stream=False, + subqueries=subqueries + ) + self.result = response + self.outcome = "success" + except BaseException as e: + self.outcome = "error" + raise e diff --git a/tests/backend/blocks/test_writerquestiontokg.py b/tests/backend/blocks/test_writerquestiontokg.py new file mode 100644 index 000000000..11c4e0911 --- /dev/null +++ b/tests/backend/blocks/test_writerquestiontokg.py @@ -0,0 +1,145 @@ +import pytest +import writer.ai +from writer.blocks.writerquestiontokg import WriterQuestionToKG +from writer.ss_types import WriterConfigurationError + +# Create a mock response object +class MockQuestion: + def __init__(self, answer="Test answer", subqueries=None): + self.answer = answer + self.subqueries = subqueries or [] + +class MockGraphs: + def question(self, graph_ids, question, stream, subqueries): + # Don't assert specific graph IDs, just verify it's a list + assert isinstance(graph_ids, list) + assert isinstance(question, str) + assert isinstance(stream, bool) + assert isinstance(subqueries, bool) + + # Return subqueries only if enabled + return MockQuestion( + answer="This is the answer", + subqueries=["subquery1", "subquery2"] if subqueries else [] + ) + +class MockClient: + def __init__(self): + self.graphs = MockGraphs() + +def mock_acquire_client(): + return MockClient() + +def test_question_to_kg(monkeypatch, session, runner): + monkeypatch.setattr(writer.ai.WriterAIManager, "acquire_client", mock_acquire_client) + + # 3. Setup: Add fake component (simulates UI input) + session.add_fake_component({ + "graphId": "123e4567-e89b-12d3-a456-426614174000", + "question": "What is the test question?", + "subqueries": "yes" + }) + + # 4. Create and run the block + block = WriterQuestionToKG("fake_id", runner, {}) + block.run() + + # 5. Verify the outcomes + assert block.outcome == "success" + assert isinstance(block.result, MockQuestion) + assert block.result.answer == "This is the answer" + assert block.result.subqueries == ["subquery1", "subquery2"] + +def test_question_to_kg_multiple_graph_ids(monkeypatch, session, runner): + monkeypatch.setattr(writer.ai.WriterAIManager, "acquire_client", mock_acquire_client) + + # 3. Setup: Add fake component (simulates UI input) + session.add_fake_component({ + "graphId": "123e4567-e89b-12d3-a456-426614174000, 123e4567-e89b-12d3-a456-426614174001", + "question": "What is the test question?", + "subqueries": "yes" + }, id="fake_id_2") + + # 4. Create and run the block + block = WriterQuestionToKG("fake_id_2", runner, {}) + block.run() + + # 5. Verify the outcomes + assert block.outcome == "success" + assert isinstance(block.result, MockQuestion) + assert block.result.answer == "This is the answer" + assert block.result.subqueries == ["subquery1", "subquery2"] + + +def test_question_to_kg_multiple_graph_ids_no_subqueries(monkeypatch, session, runner): + monkeypatch.setattr(writer.ai.WriterAIManager, "acquire_client", mock_acquire_client) + + # 3. Setup: Add fake component (simulates UI input) + session.add_fake_component({ + "graphId": "123e4567-e89b-12d3-a456-426614174000, 123e4567-e89b-12d3-a456-426614174001", + "question": "What is the test question?", + "subqueries": "no" + }) + + # 4. Create and run the block + block = WriterQuestionToKG("fake_id", runner, {}) + block.run() + + # 5. Verify the outcomes + assert block.outcome == "success" + assert isinstance(block.result, MockQuestion) + assert block.result.answer == "This is the answer" + assert block.result.subqueries == [] + + +def test_question_to_kg_missing_graph_id(monkeypatch, session, runner): + monkeypatch.setattr(writer.ai.WriterAIManager, "acquire_client", mock_acquire_client) + + # 3. Setup: Add fake component (simulates UI input) + session.add_fake_component({ + "graphId": "", + "question": "What is the test question?", + "subqueries": "yes" + }) + + # 4. Create and run the block + block = WriterQuestionToKG("fake_id", runner, {}) + with pytest.raises(WriterConfigurationError): + block.run() + +def test_question_to_kg_missing_required_fields(session, runner): + # Test missing graphId + session.add_fake_component({ + "question": "What is the test question?" + }, id="fake_id_3") + + block = WriterQuestionToKG("fake_id_3", runner, {}) + with pytest.raises(WriterConfigurationError): + block.run() + + # Test missing question + session.add_fake_component({ + "graphId": "123e4567-e89b-12d3-a456-426614174000" + }, id="fake_id_4") + block = WriterQuestionToKG("fake_id_4", runner, {}) + with pytest.raises(WriterConfigurationError): + block.run() + +def test_question_to_kg_empty_fields(session, runner): + session.add_fake_component({ + "graphId": "", + "question": "What is the test question?", + }, id="fake_id_5") + + block = WriterQuestionToKG("fake_id_5", runner, {}) + with pytest.raises(WriterConfigurationError): + block.run() + + session.add_fake_component({ + "graphId": "123e4567-e89b-12d3-a456-426614174000", + "question": "", + }, id="fake_id_6") + + block = WriterQuestionToKG("fake_id_6", runner, {}) + with pytest.raises(WriterConfigurationError): + block.run() \ No newline at end of file