Add semantic_search() method

.github/workflows/e2e.yml

@@ -15,10 +15,7 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install dependencies
-        run: |
-          sudo apt-get update
-          sudo apt-get install -y python3-pip git
-          sudo pip3 install poetry
+        run: pipx install poetry
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:

.github/workflows/unittests.yml

@@ -11,9 +11,7 @@ jobs:
         python-version: ["3.10"]
     steps:
       - uses: actions/checkout@v4
-      - run:
-          sudo apt-get update && sudo apt-get install -y python3-pip && sudo
-          pip3 install poetry
+      - run: pipx install poetry
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "yeti-python"
-version = "2.1.1"
+version = "2.2.0"
 description = "Python bindings for the Yeti API"
 authors = ["tomchop"]
 license = "Apache"

tests/api.py

@@ -112,6 +112,37 @@ def test_search_observables(self, mock_post):
             },
         )
 
+    @patch("yeti.api.requests.Session.post")
+    def test_semantic_search(self, mock_post):
+        mock_response = MagicMock()
+        mock_response.content = (
+            b'{"sections": [{"type": "entity", "results": '
+            b'[{"name": "test_entity", "semantic_score": 0.5}], "total": 1}]}'
+        )
+        mock_post.return_value = mock_response
+
+        result = self.api.semantic_search("test_query")
+        self.assertEqual(
+            result,
+            [
+                {
+                    "type": "entity",
+                    "results": [{"name": "test_entity", "semantic_score": 0.5}],
+                    "total": 1,
+                }
+            ],
+        )
+        mock_post.assert_called_with(
+            "http://fake-url/api/v2/search/semantic",
+            json={"query": "test_query", "count": 10},
+        )
+
+        result = self.api.semantic_search("test_query", count=3, root_type="dfiq")
+        mock_post.assert_called_with(
+            "http://fake-url/api/v2/search/semantic",
+            json={"query": "test_query", "count": 3, "root_type": "dfiq"},
+        )
+
     @patch("yeti.api.requests.Session.post")
     def test_search_bloom(self, mock_post):
         mock_response = MagicMock()

yeti/api.py

@@ -458,6 +458,36 @@ def search_observables(
         )
         return json.loads(response)["observables"]
 
+    def semantic_search(
+        self,
+        query: str,
+        count: int = 10,
+        root_type: str | None = None,
+    ) -> list[dict[str, Any]]:
+        """Searches Yeti for entities, indicators, or DFIQ objects related to
+        a query by meaning rather than exact text match.
+
+        Args:
+          query: A natural-language description of what to search for.
+          count: The maximum number of results to return per type (default is 10).
+          root_type: Optional. Restrict the search to a single object type:
+            "entity", "indicator", or "dfiq". Leave unset to search all types.
+
+        Returns:
+          The response from the API; a list of dicts, one per searched type,
+          each with a 'type' key, a 'results' key (a list of dicts representing
+          matched objects, each including a 'semantic_score' field where higher
+          is more similar), and a 'total' key (the number of results returned
+          for that type).
+        """
+        params: dict[str, Any] = {"query": query, "count": count}
+        if root_type:
+            params["root_type"] = root_type
+        response = self.do_request(
+            "POST", f"{self._url_root}/api/v2/search/semantic", json_data=params
+        )
+        return json.loads(response)["sections"]
+
     def search_bloom(self, values: list[str]) -> list[dict[str, Any]]:
         """Searches for a list of observable values in Yeti's bloom filters.