Merge commit '37e116e05aec5e32306b6ca5956b5515614ec9ca'

qiuosier · qiuosier · commit 15e412c0df93 · 2024-01-02T18:07:46.000-05:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,52 +1,52 @@
 repos:
-# Standard hooks
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
-    hooks:
-    -   id: check-ast
-        exclude: ^docs/
-    -   id: check-docstring-first
-        exclude: ^(docs/|tests/)
-    -   id: check-json
-    -   id: check-merge-conflict
-    -   id: check-yaml
-        exclude: feature_store_*.yaml
-        args: ['--allow-multiple-documents']
-    -   id: detect-private-key
-    -   id: end-of-file-fixer
-        exclude: '\.ipynb?$'
-    -   id: pretty-format-json
-        args: ['--autofix']
-    -   id: trailing-whitespace
-        args: [--markdown-linebreak-ext=md]
-        exclude: ^docs/
-# Black, the code formatter, natively supports pre-commit
--   repo: https://github.com/psf/black
-    rev: 23.3.0
-    hooks:
-    -   id: black
-        exclude: ^docs/
-# Regex based rst files common mistakes detector
--   repo: https://github.com/pre-commit/pygrep-hooks
-    rev: v1.10.0
-    hooks:
-    -   id: rst-backticks
-        files: ^docs/
-    -   id: rst-inline-touching-normal
-        files: ^docs/
-# Hardcoded secrets and ocids detector
--   repo: https://github.com/gitleaks/gitleaks
-    rev: v8.17.0
-    hooks:
-    -   id: gitleaks
-        exclude: .github/workflows/reusable-actions/set-dummy-conf.yml
-# Oracle copyright checker
--   repo: https://github.com/oracle-samples/oci-data-science-ai-samples/
-    rev: cbe0136f7aaffe463b31ddf3f34b0e16b4b124ff
-    hooks:
-    -   id: check-copyright
-        name: check-copyright
-        entry: .pre-commit-scripts/check-copyright.py
-        language: script
-        types_or: ['python', 'shell', 'bash']
-        exclude: ^docs/
+    # Standard hooks
+    - repo: https://github.com/pre-commit/pre-commit-hooks
+      rev: v4.4.0
+      hooks:
+          - id: check-ast
+            exclude: ^docs/
+          - id: check-docstring-first
+            exclude: ^(docs/|tests/)
+          - id: check-json
+          - id: check-merge-conflict
+          - id: check-yaml
+            exclude: feature_store_*.yaml
+            args: ["--allow-multiple-documents"]
+          - id: detect-private-key
+          - id: end-of-file-fixer
+            exclude: '\.ipynb?$'
+          - id: pretty-format-json
+            args: ["--autofix"]
+          - id: trailing-whitespace
+            args: [--markdown-linebreak-ext=md]
+            exclude: ^docs/
+    # Black, the code formatter, natively supports pre-commit
+    - repo: https://github.com/psf/black
+      rev: 23.3.0
+      hooks:
+          - id: black
+            exclude: ^docs/
+    # Regex based rst files common mistakes detector
+    - repo: https://github.com/pre-commit/pygrep-hooks
+      rev: v1.10.0
+      hooks:
+          - id: rst-backticks
+            files: ^docs/
+          - id: rst-inline-touching-normal
+            files: ^docs/
+    # Hardcoded secrets and ocids detector
+    - repo: https://github.com/gitleaks/gitleaks
+      rev: v8.17.0
+      hooks:
+          - id: gitleaks
+            exclude: .github/workflows/reusable-actions/set-dummy-conf.yml
+    # Oracle copyright checker
+    - repo: https://github.com/oracle-samples/oci-data-science-ai-samples/
+      rev: 1bc5270a443b791c62f634233c0f4966dfcc0dd6
+      hooks:
+          - id: check-copyright
+            name: check-copyright
+            entry: .pre-commit-scripts/check-copyright.py
+            language: script
+            types_or: ["python", "shell", "bash"]
+            exclude: ^docs/
diff --git a/ads/llm/chain.py b/ads/llm/chain.py
@@ -261,12 +261,8 @@ def load(cls, chain_dict: dict, **kwargs) -> "GuardrailSequence":
         from ads.llm.serialize import load
 
         chain_spec = chain_dict[SPEC_CHAIN]
-        chain = cls()
-        for config in chain_spec:
-            step = load(config, **kwargs)
-            # Chain the step
-            chain |= step
-        return chain
+        steps = [load(config, **kwargs) for config in chain_spec]
+        return cls(*steps)
 
     def __str__(self) -> str:
         return "\n".join([str(step.__class__) for step in self.steps])
diff --git a/tests/unitary/with_extras/langchain/test_guardrails.py b/tests/unitary/with_extras/langchain/test_guardrails.py
@@ -175,10 +175,6 @@ def test_fn(chain: GuardrailSequence):
 
         self.assert_before_and_after_serialization(test_fn, chain)
 
-    def test_empty_sequence(self):
-        """Tests empty sequence."""
-        seq = GuardrailSequence()
-        self.assertEqual(seq.steps, [])
 
     def test_save_to_file(self):
         """Tests saving to file."""
diff --git a/tests/unitary/with_extras/langchain/test_serialization.py b/tests/unitary/with_extras/langchain/test_serialization.py
@@ -40,39 +40,6 @@ def setUp(self) -> None:
     GEN_AI_KWARGS = {"service_endpoint": "https://endpoint.oraclecloud.com"}
     ENDPOINT = "https://modeldeployment.customer-oci.com/ocid/predict"
 
-    EXPECTED_LLM_CHAIN_WITH_COHERE = {
-        "memory": None,
-        "verbose": True,
-        "tags": None,
-        "metadata": None,
-        "prompt": {
-            "input_variables": ["subject"],
-            "input_types": {},
-            "output_parser": None,
-            "partial_variables": {},
-            "template": "Tell me a joke about {subject}",
-            "template_format": "f-string",
-            "validate_template": False,
-            "_type": "prompt",
-        },
-        "llm": {
-            "model": None,
-            "max_tokens": 256,
-            "temperature": 0.75,
-            "k": 0,
-            "p": 1,
-            "frequency_penalty": 0.0,
-            "presence_penalty": 0.0,
-            "truncate": None,
-            "_type": "cohere",
-        },
-        "output_key": "text",
-        "output_parser": {"_type": "default"},
-        "return_final_only": True,
-        "llm_kwargs": {},
-        "_type": "llm_chain",
-    }
-
     EXPECTED_LLM_CHAIN_WITH_OCI_MD = {
         "lc": 1,
         "type": "constructor",
@@ -173,7 +140,23 @@ def test_llm_chain_serialization_with_cohere(self):
         template = PromptTemplate.from_template(self.PROMPT_TEMPLATE)
         llm_chain = LLMChain(prompt=template, llm=llm, verbose=True)
         serialized = dump(llm_chain)
-        self.assertEqual(serialized, self.EXPECTED_LLM_CHAIN_WITH_COHERE)
+
+        # Check the serialized chain
+        self.assertTrue(serialized.get("verbose"))
+        self.assertEqual(serialized.get("_type"), "llm_chain")
+
+        # Check the serialized prompt template
+        serialized_prompt = serialized.get("prompt")
+        self.assertIsInstance(serialized_prompt, dict)
+        self.assertEqual(serialized_prompt.get("_type"), "prompt")
+        self.assertEqual(set(serialized_prompt.get("input_variables")), {"subject"})
+        self.assertEqual(serialized_prompt.get("template"), self.PROMPT_TEMPLATE)
+
+        # Check the serialized LLM
+        serialized_llm = serialized.get("llm")
+        self.assertIsInstance(serialized_llm, dict)
+        self.assertEqual(serialized_llm.get("_type"), "cohere")
+
         llm_chain = load(serialized)
         self.assertIsInstance(llm_chain, LLMChain)
         self.assertIsInstance(llm_chain.prompt, PromptTemplate)
@@ -237,21 +220,37 @@ def test_runnable_sequence_serialization(self):
 
         chain = map_input | template | llm
         serialized = dump(chain)
-        # Do not check the ID fields.
-        expected = deepcopy(self.EXPECTED_RUNNABLE_SEQUENCE)
-        expected["id"] = serialized["id"]
-        expected["kwargs"]["first"]["id"] = serialized["kwargs"]["first"]["id"]
-        expected["kwargs"]["first"]["kwargs"]["steps"]["text"]["id"] = serialized[
-            "kwargs"
-        ]["first"]["kwargs"]["steps"]["text"]["id"]
-        expected["kwargs"]["middle"][0]["id"] = serialized["kwargs"]["middle"][0]["id"]
-        self.assertEqual(serialized, expected)
+
+        self.assertEqual(serialized.get("type"), "constructor")
+        self.assertNotIn("_type", serialized)
+
+        kwargs = serialized.get("kwargs")
+        self.assertIsInstance(kwargs, dict)
+
+        element_1 = kwargs.get("first")
+        self.assertEqual(element_1.get("_type"), "RunnableParallel")
+        step = element_1.get("kwargs").get("steps").get("text")
+        self.assertEqual(step.get("id")[-1], "RunnablePassthrough")
+
+        element_2 = kwargs.get("middle")[0]
+        self.assertNotIn("_type", element_2)
+        self.assertEqual(element_2.get("kwargs").get("template"), self.PROMPT_TEMPLATE)
+        self.assertEqual(element_2.get("kwargs").get("input_variables"), ["subject"])
+
+        element_3 = kwargs.get("last")
+        self.assertNotIn("_type", element_3)
+        self.assertEqual(element_3.get("id"), ["ads", "llm", "ModelDeploymentTGI"])
+        self.assertEqual(
+            element_3.get("kwargs"),
+            {"endpoint": "https://modeldeployment.customer-oci.com/ocid/predict"},
+        )
+
         chain = load(serialized)
         self.assertEqual(len(chain.steps), 3)
         self.assertIsInstance(chain.steps[0], RunnableParallel)
         self.assertEqual(
-            chain.steps[0].dict(),
-            {"steps": {"text": {"input_type": None, "func": None, "afunc": None}}},
+            list(chain.steps[0].dict().get("steps").keys()),
+            ["text"],
         )
         self.assertIsInstance(chain.steps[1], PromptTemplate)
         self.assertIsInstance(chain.steps[2], ModelDeploymentTGI)
