diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index e2fef04..441904f 100644 --- a/.github/workflows/e2e.yml +++ b/.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: diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml index f4cae35..94a5093 100644 --- a/.github/workflows/unittests.yml +++ b/.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: diff --git a/pyproject.toml b/pyproject.toml index 1d45731..b7d2217 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/api.py b/tests/api.py index d28e199..4919a21 100644 --- a/tests/api.py +++ b/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() diff --git a/yeti/api.py b/yeti/api.py index dc0e82e..79ad183 100644 --- a/yeti/api.py +++ b/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.