From 0a75b06d3b41a8f666abc7e9018c301d46730e17 Mon Sep 17 00:00:00 2001
From: Estelle Poulin <dev@inspiredby.es>
Date: Tue, 7 Nov 2023 10:41:45 -0500
Subject: [PATCH] News endpoint

---
 README.md            |  4 +++-
 kaginawa/__init__.py |  2 +-
 kaginawa/client.py   | 27 +++++++++++++++++++++++++--
 kaginawa/models.py   |  2 +-
 test/test_models.py  |  6 +++---
 5 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index b8c3c58..4fcb72d 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,9 @@ for reference in response.references:
     print(reference.snippet)
     print(reference.url)
 
-response: KaginawaEnrichWebResponse = client.enrich_web(query="Best fermented hot sauce")
+response: KaginawaEnrichResponse = client.enrich_web(query="Best fermented hot sauce")
+# or 
+response: KaginawaEnrichResponse = client.enrich_news(query="Is Oliver Tree okay?")
 
 for result in response.results:
     print(result.rank)
diff --git a/kaginawa/__init__.py b/kaginawa/__init__.py
index a13031a..ba88f37 100644
--- a/kaginawa/__init__.py
+++ b/kaginawa/__init__.py
@@ -1,3 +1,3 @@
 from pkg_resources import parse_version
 
-__version__ = parse_version("0.0.6")
+__version__ = parse_version("0.0.7")
diff --git a/kaginawa/client.py b/kaginawa/client.py
index 07a2a46..e5b7c7c 100644
--- a/kaginawa/client.py
+++ b/kaginawa/client.py
@@ -5,7 +5,7 @@
 
 from kaginawa.exceptions import KaginawaError
 from kaginawa.models import (
-    KaginawaEnrichWebResponse,
+    KaginawaEnrichResponse,
     KaginawaFastGPTResponse,
     KaginawaSummarizationResponse,
 )
@@ -104,7 +104,30 @@ def enrich_web(self, query: str):
         except requests.RequestException as e:
             raise KaginawaError("Error calling /v0/enrich/web") from e
 
-        return KaginawaEnrichWebResponse.from_raw(raw_response)
+        return KaginawaEnrichResponse.from_raw(raw_response).results[0]
+
+    def enrich_news(self, query: str):
+        """Query the Teclis index for relevant web results for a given query.
+
+        Parameters:
+          query (str): The search query.
+
+        Returns:
+          KaginawaEnrichWebResponse: A model representing the response.
+        """
+
+        try:
+            res = self.session.get(
+                f"{self.api_base}/v0/enrich/news",
+                params={"q": query},
+            )
+            res.raise_for_status()
+
+            raw_response = res.json()
+        except requests.RequestException as e:
+            raise KaginawaError("Error calling /v0/enrich/web") from e
+
+        return KaginawaEnrichResponse.from_raw(raw_response).results[0]
 
     def summarize(
         self,
diff --git a/kaginawa/models.py b/kaginawa/models.py
index c5b4509..ed789ec 100644
--- a/kaginawa/models.py
+++ b/kaginawa/models.py
@@ -146,7 +146,7 @@ def from_raw(cls, raw_result: Dict[str, Any]):
 
 
 @dataclass
-class KaginawaEnrichWebResponse(KaginawaResponse):
+class KaginawaEnrichResponse(KaginawaResponse):
     results: List[KaginawaSearchResult]
 
     @classmethod
diff --git a/test/test_models.py b/test/test_models.py
index 134be10..de87e55 100644
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -1,7 +1,7 @@
 from datetime import datetime, timedelta
 
 from kaginawa.models import (
-    KaginawaEnrichWebResponse,
+    KaginawaEnrichResponse,
     KaginawaFastGPTResponse,
     KaginawaResponse,
     KaginawaSearchResult,
@@ -105,7 +105,7 @@ def test_kagi_fastgpt_response_no_references(self):
         assert res.tokens == tokens
         assert len(res.references) == 0
 
-    def test_kagi_enrich_web_response(self):
+    def test_kagi_enrich_response(self):
         enrich_web_response = {
             "meta": {
                 "id": (id := "cbad1ea2-ddef-42e7-af75-81bce2c64924"),
@@ -127,7 +127,7 @@ def test_kagi_enrich_web_response(self):
             ),
         }
 
-        res = KaginawaEnrichWebResponse.from_raw(enrich_web_response)
+        res = KaginawaEnrichResponse.from_raw(enrich_web_response)
 
         assert res.id == id
         assert res.node == node