참고소스 수정본

참고/guardrails-main/tests/integration_tests/integrations/langchain/test_guard_runnable.py

@@ -0,0 +1,145 @@
+from typing import Any, Optional
+import io
+import sys
+
+import pytest
+from pydantic import PrivateAttr
+
+from guardrails.guard import Guard
+from guardrails.integrations.langchain.guard_runnable import GuardRunnable
+from guardrails.errors import ValidationError
+from guardrails.classes import ValidationOutcome
+
+from tests.integration_tests.test_assets.validators import ReadingTime, RegexMatch
+
+
+@pytest.fixture
+def guard_runnable():
+    return GuardRunnable(
+        Guard().use(
+            RegexMatch("Ice cream", match_type="search", on_fail="refrain"),
+            ReadingTime(0.05, on_fail="noop"),
+            on="output",
+        )
+    )
+
+
+@pytest.mark.parametrize(
+    "output,throws",
+    [
+        ("Ice cream is frozen.", False),
+        ("Ice cream is a frozen dairy product that is consumed in many places.", True),
+        ("This response isn't relevant.", True),
+    ],
+)
+def test_guard_as_runnable(guard_runnable: GuardRunnable, output: str, throws: bool):
+    from langchain_core.language_models import LanguageModelInput
+    from langchain_core.messages import AIMessage, BaseMessage
+    from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_core.runnables import Runnable, RunnableConfig
+
+    class MockModel(Runnable):
+        def invoke(
+            self,
+            input: LanguageModelInput,
+            config: Optional[RunnableConfig] = None,
+            **kwargs: Any,
+        ) -> BaseMessage:
+            return AIMessage(content=output)
+
+    prompt = ChatPromptTemplate.from_template("ELIF: {topic}")
+    model = MockModel()
+    output_parser = StrOutputParser()
+
+    chain = prompt | model | guard_runnable | output_parser
+
+    topic = "ice cream"
+    if throws:
+        with pytest.raises(ValidationError) as exc_info:
+            chain.invoke({"topic": topic})
+
+        assert str(exc_info.value) == (
+            "The response from the LLM failed validation!"
+            "See `guard.history` for more details."
+        )
+
+        assert guard_runnable.guard.history.last.status == "fail"
+        assert guard_runnable.guard.history.last.status == "fail"
+
+    else:
+        result = chain.invoke({"topic": topic})
+
+        assert result == output
+
+
+def test_guard_runnable_with_callback_config(guard_runnable):
+    from langchain_core.callbacks import CallbackManager
+    from langchain_core.tracers import ConsoleCallbackHandler
+    from langchain_core.runnables import RunnableConfig
+
+    console_handler = ConsoleCallbackHandler()
+    callback_manager = CallbackManager([console_handler])
+    config_with_callbacks = RunnableConfig(callbacks=callback_manager)
+
+    captured_output = io.StringIO()
+    sys.stdout = captured_output
+    guard_runnable.invoke("Ice cream is sweet", config=config_with_callbacks)
+    sys.stdout = sys.__stdout__
+
+    assert "Ice cream is sweet" in captured_output.getvalue()
+
+
+@pytest.mark.parametrize(
+    "succeed_on_attempt, max_retries, expected_attempts, expected_result",
+    [
+        (1, 2, 1, "Succeeded on attempt 1"),
+        (2, 2, 1, "Failed attempt 1"),
+        (2, None, 1, "Failed attempt 1"),
+        (2, 0, 1, "Failed attempt 1"),
+    ],
+)
+def test_guard_runnable_max_retries(
+    succeed_on_attempt, max_retries, expected_attempts, expected_result
+):
+    from langchain_core.runnables import RunnableConfig
+
+    class CountingGuard(Guard):
+        _attempt_count: int = PrivateAttr(default=0)
+        _succeed_on_attempt: int = PrivateAttr()
+
+        def __init__(self, succeed_on_attempt: int, **kwargs):
+            super().__init__(**kwargs)
+            self._succeed_on_attempt = succeed_on_attempt
+
+        def validate(self, value):
+            self._attempt_count += 1
+            if self._attempt_count >= self._succeed_on_attempt:
+                return ValidationOutcome(
+                    call_id="0",  # type: ignore
+                    raw_llm_output=value,
+                    validated_output=f"Succeeded on attempt {self._attempt_count}",
+                    validation_passed=True,
+                )
+            raise ValidationError(f"Failed attempt {self._attempt_count}")
+
+        @property
+        def attempt_count(self):
+            return self._attempt_count
+
+    guard = CountingGuard(succeed_on_attempt)
+    runnable = GuardRunnable(guard)
+
+    config = (
+        RunnableConfig(max_retries=max_retries) if max_retries is not None else None
+    )
+
+    if "Failed" in expected_result:
+        with pytest.raises(ValidationError) as exc_info:
+            runnable.invoke("test input", config=config)
+        assert expected_result in str(exc_info.value)
+    else:
+        result = runnable.invoke("test input", config=config)
+        assert result == expected_result
+
+    assert guard.attempt_count == expected_attempts

참고/guardrails-main/tests/integration_tests/integrations/langchain/test_validator_runnable.py

@@ -0,0 +1,95 @@
+from typing import Any, Optional
+import io
+import sys
+
+
+import pytest
+
+from guardrails.errors import ValidationError
+from tests.integration_tests.test_assets.validators import ReadingTime, RegexMatch
+
+
+@pytest.mark.parametrize(
+    "output,throws,expected_error",
+    [
+        ("Ice cream is frozen.", False, None),
+        (
+            "Ice cream is a frozen dairy product that is consumed in many places.",
+            True,
+            "String should be readable within 0.05 minutes.",
+        ),
+        ("This response isn't relevant.", True, "Result must match Ice cream"),
+    ],
+)
+def test_validator_runnable(output: str, throws: bool, expected_error: Optional[str]):
+    from langchain_core.language_models import LanguageModelInput
+    from langchain_core.messages import AIMessage, BaseMessage
+    from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_core.runnables import Runnable, RunnableConfig
+
+    class MockModel(Runnable):
+        def invoke(
+            self,
+            input: LanguageModelInput,
+            config: Optional[RunnableConfig] = None,
+            **kwargs: Any,
+        ) -> BaseMessage:
+            return AIMessage(content=output)
+
+    prompt = ChatPromptTemplate.from_template("ELIF: {topic}")
+    model = MockModel()
+    regex_match = RegexMatch(
+        "Ice cream", match_type="search", on_fail="refrain"
+    ).to_runnable()
+    reading_time = ReadingTime(0.05, on_fail="refrain").to_runnable()
+
+    output_parser = StrOutputParser()
+
+    chain = prompt | model | regex_match | reading_time | output_parser
+
+    topic = "ice cream"
+    if throws:
+        with pytest.raises(ValidationError) as exc_info:
+            chain.invoke({"topic": topic})
+
+        assert str(exc_info.value) == (
+            f"The response from the LLM failed validation! {expected_error}"
+        )
+
+    else:
+        result = chain.invoke({"topic": topic})
+
+        assert result == output
+
+
+def test_validator_runnable_with_callback_config():
+    from langchain_core.callbacks import CallbackManager
+    from langchain_core.tracers import ConsoleCallbackHandler
+    from langchain_core.runnables import RunnableConfig
+
+    console_handler = ConsoleCallbackHandler()
+    callback_manager = CallbackManager([console_handler])
+    config_with_callbacks = RunnableConfig(callbacks=callback_manager)
+
+    regex_match = RegexMatch(
+        "Ice cream", match_type="search", on_fail="exception"
+    ).to_runnable()
+
+    captured_output = io.StringIO()
+    sys.stdout = captured_output
+
+    result = regex_match.invoke("Ice cream is delicious.", config=config_with_callbacks)
+    assert result == "Ice cream is delicious."
+
+    with pytest.raises(ValidationError) as exc_info:
+        regex_match.invoke("Chocolate is delicious.", config=config_with_callbacks)
+
+    assert "The response from the LLM failed validation!" in str(exc_info.value)
+    assert "Result must match Ice cream" in str(exc_info.value)
+
+    sys.stdout = sys.__stdout__
+
+    console_output = captured_output.getvalue()
+    assert "Ice cream is delicious." in console_output
+    assert "Chocolate is delicious." in console_output

참고/guardrails-main/tests/integration_tests/integrations/llama_index/test_guardrails_chat_engine.py

@@ -0,0 +1,67 @@
+import pytest
+from guardrails import Guard
+from typing import List, Optional
+from tests.integration_tests.test_assets.validators import RegexMatch
+
+pytest.importorskip("llama_index")
+
+from llama_index.core.chat_engine.types import (  # noqa
+    BaseChatEngine,  # noqa
+    AgentChatResponse,  # noqa
+    StreamingAgentChatResponse,  # noqa
+)  # noqa
+from llama_index.core.base.llms.types import ChatMessage  # noqa
+from guardrails.integrations.llama_index import GuardrailsChatEngine  # noqa
+
+
+class MockChatEngine(BaseChatEngine):
+    def chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> AgentChatResponse:
+        return AgentChatResponse(response="Mock response")
+
+    async def achat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ) -> AgentChatResponse:
+        return AgentChatResponse(response="Mock async chat response")
+
+    def stream_chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ):
+        return StreamingAgentChatResponse(response="Mock stream chat response")
+
+    async def astream_chat(
+        self, message: str, chat_history: Optional[List[ChatMessage]] = None
+    ):
+        return StreamingAgentChatResponse(response="Mock async stream chat response")
+
+    @property
+    def chat_history(self) -> List[ChatMessage]:
+        return []
+
+    def reset(self):
+        pass
+
+
+pytest.importorskip("llama_index")
+
+
+@pytest.fixture
+def guard():
+    return Guard().use(RegexMatch("Mock response", match_type="search"))
+
+
+class TestGuardrailsChatEngine:
+    def test_guardrails_engine_init(self, guard):
+        engine = MockChatEngine()
+        guardrails_engine = GuardrailsChatEngine(engine, guard)
+        assert isinstance(guardrails_engine, GuardrailsChatEngine)
+        assert guardrails_engine.guard == guard
+
+    def test_guardrails_engine_chat(self, guard):
+        engine = MockChatEngine()
+        guardrails_engine = GuardrailsChatEngine(engine, guard)
+
+        result = guardrails_engine.chat("Mock response")
+        assert isinstance(result, AgentChatResponse)
+        assert result.response == "Mock response"

참고/guardrails-main/tests/integration_tests/integrations/llama_index/test_guardrails_query_engine.py

@@ -0,0 +1,63 @@
+import pytest
+from guardrails import Guard
+from guardrails.errors import ValidationError
+from typing import Optional
+from tests.integration_tests.test_assets.validators import RegexMatch
+
+pytest.importorskip("llama_index")
+
+from llama_index.core.query_engine import BaseQueryEngine  # noqa
+from llama_index.core.schema import QueryBundle  # noqa
+from llama_index.core.base.response.schema import Response  # noqa
+from llama_index.core.prompts.mixin import PromptMixinType  # noqa
+from llama_index.core.callbacks import CallbackManager  # noqa
+
+
+class MockQueryEngine(BaseQueryEngine):
+    def __init__(self, callback_manager: Optional[CallbackManager] = None):
+        super().__init__(callback_manager)
+
+    def _query(self, query_bundle: QueryBundle) -> Response:
+        return Response(response="Mock response")
+
+    async def _aquery(self, query_bundle: QueryBundle) -> Response:
+        return Response(response="Mock async query response")
+
+    def _get_prompt_modules(self) -> PromptMixinType:
+        return {}
+
+
+@pytest.fixture
+def guard():
+    return Guard().use(RegexMatch("Mock response", match_type="search"))
+
+
+class TestGuardrailsQueryEngine:
+    def test_guardrails_engine_init(self, guard):
+        from guardrails.integrations.llama_index import GuardrailsQueryEngine
+
+        engine = MockQueryEngine()
+        guardrails_engine = GuardrailsQueryEngine(engine, guard)
+        assert isinstance(guardrails_engine, GuardrailsQueryEngine)
+        assert guardrails_engine.guard == guard
+
+    def test_guardrails_engine_query(self, guard):
+        from guardrails.integrations.llama_index import GuardrailsQueryEngine
+
+        engine = MockQueryEngine()
+        guardrails_engine = GuardrailsQueryEngine(engine, guard)
+
+        result = guardrails_engine._query(QueryBundle(query_str="Mock response"))
+        assert isinstance(result, Response)
+        assert result.response == "Mock response"
+
+    def test_guardrails_engine_query_validation_failure(self, guard):
+        from guardrails.integrations.llama_index import GuardrailsQueryEngine
+
+        engine = MockQueryEngine()
+        guardrails_engine = GuardrailsQueryEngine(engine, guard)
+
+        engine._query = lambda _: Response(response="Invalid response")
+
+        with pytest.raises(ValidationError, match="Validation failed"):
+            guardrails_engine._query(QueryBundle(query_str="Invalid query"))