Skip to content

Instantly share code, notes, and snippets.

@vcidst
Created October 4, 2024 20:42
Show Gist options
  • Select an option

  • Save vcidst/9adaa3fd0a3e0a98115d4e0c171f09fe to your computer and use it in GitHub Desktop.

Select an option

Save vcidst/9adaa3fd0a3e0a98115d4e0c171f09fe to your computer and use it in GitHub Desktop.
kapa ai rasa custom information retrieval
from typing import Text, Any, Dict
import requests
from rasa.utils.endpoints import EndpointConfig
from rasa.core.information_retrieval import (
SearchResult,
SearchResultList,
InformationRetrieval,
)
class Kapa_AI(InformationRetrieval):
def connect(
self,
config: EndpointConfig,
) -> None:
"""Initialise Kapa AI API Parameters."""
self.api_key = "kapa-ai-api-key"
project_id = "kapa-project-id"
self.base_url = f"https://api.kapa.ai/query/v1/projects/{project_id}"
self.headers = {
"Authorization": f"Bearer {self.api_key}",
}
def _create_search_results(self, response_json):
"""Helper method to create search results from API response."""
hits = SearchResultList()
for result in response_json:
hits.results.append(
SearchResult(
text=result.get("content"),
metadata={
"source_url": result.get("source_url"),
"title": result.get("title"),
"source_type": result.get("source_type"),
},
)
)
return hits
async def search(
self, query: Text, tracker_state: Dict[str, Any], threshold: float = 0.0
) -> SearchResultList:
"""Search Kapa AI for relevant documents."""
url = self.base_url + "/search"
data = {"query": query, "num_results": 4, "include_source_names": ["string"]}
response = requests.post(
url,
headers=self.headers,
json=data
)
return self._create_search_results(response.json())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment