mirror of
https://github.com/khoj-ai/khoj.git
synced 2025-02-17 08:04:21 +00:00
Setup Pytest fixture for conversation processor to test chat API
- Index markdown test data as knowledge base. As easier to get good markdown content (vs org) - Setup markdown_content_config, processor_config and chat_client to test chat API
This commit is contained in:
parent
3f719c9e17
commit
b6d63137f1
1 changed files with 62 additions and 2 deletions
|
@ -1,4 +1,5 @@
|
|||
# External Packages
|
||||
import os
|
||||
from copy import deepcopy
|
||||
from fastapi.testclient import TestClient
|
||||
from pathlib import Path
|
||||
|
@ -6,11 +7,14 @@ import pytest
|
|||
|
||||
# Internal Packages
|
||||
from khoj.main import app
|
||||
from khoj.configure import configure_routes, configure_search_types
|
||||
from khoj.configure import configure_processor, configure_routes, configure_search_types
|
||||
from khoj.processor.markdown.markdown_to_jsonl import MarkdownToJsonl
|
||||
from khoj.search_type import image_search, text_search
|
||||
from khoj.utils.helpers import resolve_absolute_path
|
||||
from khoj.utils.rawconfig import (
|
||||
ContentConfig,
|
||||
ConversationProcessorConfig,
|
||||
ProcessorConfig,
|
||||
TextContentConfig,
|
||||
ImageContentConfig,
|
||||
SearchConfig,
|
||||
|
@ -94,7 +98,63 @@ def content_config(tmp_path_factory, search_config: SearchConfig):
|
|||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client(content_config: ContentConfig, search_config: SearchConfig):
|
||||
def md_content_config(tmp_path_factory):
|
||||
content_dir = tmp_path_factory.mktemp("content")
|
||||
|
||||
# Generate Embeddings for Markdown Content
|
||||
content_config = ContentConfig()
|
||||
content_config.markdown = TextContentConfig(
|
||||
input_files=None,
|
||||
input_filter=["tests/data/markdown/*.md"],
|
||||
compressed_jsonl=content_dir.joinpath("markdown.jsonl"),
|
||||
embeddings_file=content_dir.joinpath("markdown_embeddings.pt"),
|
||||
)
|
||||
|
||||
return content_config
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def processor_config(tmp_path_factory):
|
||||
openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
processor_dir = tmp_path_factory.mktemp("processor")
|
||||
|
||||
# The conversation processor is the only configured processor
|
||||
# It needs an OpenAI API key to work.
|
||||
if not openai_api_key:
|
||||
return
|
||||
|
||||
# Setup conversation processor, if OpenAI API key is set
|
||||
processor_config = ProcessorConfig()
|
||||
processor_config.conversation = ConversationProcessorConfig(
|
||||
openai_api_key=openai_api_key,
|
||||
conversation_logfile=processor_dir.joinpath("conversation_logs.json"),
|
||||
)
|
||||
|
||||
return processor_config
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def chat_client(md_content_config: ContentConfig, search_config: SearchConfig, processor_config: ProcessorConfig):
|
||||
# Initialize app state
|
||||
state.config.content_type = md_content_config
|
||||
state.config.search_type = search_config
|
||||
state.SearchType = configure_search_types(state.config)
|
||||
|
||||
# Index Markdown Content for Search
|
||||
filters = [DateFilter(), WordFilter(), FileFilter()]
|
||||
state.model.markdown_search = text_search.setup(
|
||||
MarkdownToJsonl, md_content_config.markdown, search_config.asymmetric, regenerate=False, filters=filters
|
||||
)
|
||||
|
||||
# Initialize Processor from Config
|
||||
state.processor_config = configure_processor(processor_config)
|
||||
|
||||
configure_routes(app)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def client(content_config: ContentConfig, search_config: SearchConfig, processor_config: ProcessorConfig):
|
||||
state.config.content_type = content_config
|
||||
state.config.search_type = search_config
|
||||
state.SearchType = configure_search_types(state.config)
|
||||
|
|
Loading…
Add table
Reference in a new issue