Add search task type and provider only returning sources

appinfo/info.xml

@@ -48,6 +48,7 @@ Refer to the [Context Chat Backend's readme](https://github.com/nextcloud/contex
 	</background-jobs>
 	<commands>
 		<command>OCA\ContextChat\Command\Prompt</command>
+		<command>OCA\ContextChat\Command\Search</command>
 		<command>OCA\ContextChat\Command\ScanFiles</command>
 		<command>OCA\ContextChat\Command\Statistics</command>
 	</commands>

lib/AppInfo/Application.php

@@ -12,6 +12,8 @@
 use OCA\ContextChat\Listener\ShareListener;
 use OCA\ContextChat\Listener\UserDeletedListener;
 use OCA\ContextChat\TaskProcessing\ContextChatProvider;
+use OCA\ContextChat\TaskProcessing\ContextChatSearchProvider;
+use OCA\ContextChat\TaskProcessing\ContextChatSearchTaskType;
 use OCA\ContextChat\TaskProcessing\ContextChatTaskType;
 use OCP\App\Events\AppDisableEvent;
 use OCP\AppFramework\App;
@@ -82,6 +84,8 @@ public function register(IRegistrationContext $context): void {
 		$context->registerEventListener(ShareDeletedEvent::class, ShareListener::class);
 		$context->registerTaskProcessingTaskType(ContextChatTaskType::class);
 		$context->registerTaskProcessingProvider(ContextChatProvider::class);
+		$context->registerTaskProcessingTaskType(ContextChatSearchTaskType::class);
+		$context->registerTaskProcessingProvider(ContextChatSearchProvider::class);
 	}
 
 	public function boot(IBootContext $context): void {

lib/Command/Search.php

@@ -0,0 +1,84 @@
+<?php
+
+/**
+ * SPDX-FileCopyrightText: 2023 Nextcloud GmbH and Nextcloud contributors
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+namespace OCA\ContextChat\Command;
+
+use OCA\ContextChat\TaskProcessing\ContextChatSearchTaskType;
+use OCA\ContextChat\Type\ScopeType;
+use OCP\TaskProcessing\IManager;
+use OCP\TaskProcessing\Task;
+use Symfony\Component\Console\Command\Command;
+use Symfony\Component\Console\Input\InputArgument;
+use Symfony\Component\Console\Input\InputInterface;
+use Symfony\Component\Console\Input\InputOption;
+use Symfony\Component\Console\Output\OutputInterface;
+
+class Search extends Command {
+
+	public function __construct(
+		private IManager $taskProcessingManager,
+	) {
+		parent::__construct();
+	}
+
+	protected function configure() {
+		$this->setName('context_chat:search')
+			->setDescription('Search with Nextcloud Assistant Context Chat')
+			->addArgument(
+				'uid',
+				InputArgument::REQUIRED,
+				'The ID of the user to search the documents of'
+			)
+			->addArgument(
+				'prompt',
+				InputArgument::REQUIRED,
+				'The prompt'
+			)
+			->addOption(
+				'context-providers',
+				null,
+				InputOption::VALUE_REQUIRED,
+				'Context providers to use (as a comma-separated list without brackets)',
+			);
+	}
+
+	protected function execute(InputInterface $input, OutputInterface $output) {
+		$userId = $input->getArgument('uid');
+		$prompt = $input->getArgument('prompt');
+		$contextProviders = $input->getOption('context-providers');
+
+		if (!empty($contextProviders)) {
+			$contextProviders = preg_replace('/\s*,+\s*/', ',', $contextProviders);
+			$contextProvidersArray = array_filter(explode(',', $contextProviders), fn ($source) => !empty($source));
+			$task = new Task(ContextChatSearchTaskType::ID, [
+				'prompt' => $prompt,
+				'scopeType' => ScopeType::PROVIDER,
+				'scopeList' => $contextProvidersArray,
+				'scopeListMeta' => '',
+			], 'context_chat', $userId);
+		} else {
+			$task = new Task(ContextChatSearchTaskType::ID, [
+				'prompt' => $prompt,
+				'scopeType' => ScopeType::NONE,
+				'scopeList' => [],
+				'scopeListMeta' => '',
+			], 'context_chat', $userId);
+		}
+
+		$this->taskProcessingManager->scheduleTask($task);
+		while (!in_array(($task = $this->taskProcessingManager->getTask($task->getId()))->getStatus(), [Task::STATUS_FAILED, Task::STATUS_SUCCESSFUL], true)) {
+			sleep(1);
+		}
+		if ($task->getStatus() === Task::STATUS_SUCCESSFUL) {
+			$output->writeln(var_export($task->getOutput(), true));
+			return 0;
+		} else {
+			$output->writeln('<error>' . $task->getErrorMessage() . '</error>');
+			return 1;
+		}
+	}
+}

lib/Service/LangRopeService.php

@@ -334,6 +334,30 @@ public function query(string $userId, string $prompt, bool $useContext = true, ?
 		return $this->requestToExApp('/query', 'POST', $params);
 	}
 
+	/**
+	 * @param string $userId
+	 * @param string $prompt
+	 * @param ?string $scopeType
+	 * @param ?array<string> $scopeList
+	 * @param int|null $limit
+	 * @return array
+	 */
+	public function docSearch(string $userId, string $prompt, ?string $scopeType = null, ?array $scopeList = null, ?int $limit = null): array {
+		$params = [
+			'query' => $prompt,
+			'userId' => $userId,
+		];
+		if ($scopeType !== null && $scopeList !== null) {
+			$params['scopeType'] = $scopeType;
+			$params['scopeList'] = $scopeList;
+		}
+		if ($limit !== null) {
+			$params['ctxLimit'] = $limit;
+		}
+
+		return $this->requestToExApp('/docSearch', 'POST', $params);
+	}
+
 	/**
 	 * @param string $userId
 	 * @param string $prompt

lib/TaskProcessing/ContextChatSearchProvider.php

@@ -0,0 +1,206 @@
+<?php
+
+declare(strict_types=1);
+
+/**
+ * SPDX-FileCopyrightText: 2024 Nextcloud GmbH and Nextcloud contributors
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+namespace OCA\ContextChat\TaskProcessing;
+
+use OCA\ContextChat\AppInfo\Application;
+use OCA\ContextChat\Logger;
+use OCA\ContextChat\Service\LangRopeService;
+use OCA\ContextChat\Service\MetadataService;
+use OCA\ContextChat\Type\ScopeType;
+use OCP\IL10N;
+use OCP\TaskProcessing\ISynchronousProvider;
+
+class ContextChatSearchProvider implements ISynchronousProvider {
+
+	public function __construct(
+		private LangRopeService $langRopeService,
+		private IL10N $l10n,
+		private Logger $logger,
+		private MetadataService $metadataService,
+	) {
+	}
+
+	public function getId(): string {
+		return Application::APP_ID . '-context_chat_search';
+	}
+
+	public function getName(): string {
+		return $this->l10n->t('Nextcloud Assistant Context Chat Search Provider');
+	}
+
+	public function getTaskTypeId(): string {
+		return ContextChatSearchTaskType::ID;
+	}
+
+	public function getExpectedRuntime(): int {
+		return 120;
+	}
+
+	public function getInputShapeEnumValues(): array {
+		return [];
+	}
+
+	public function getInputShapeDefaults(): array {
+		return [
+			'limit' => 10,
+		];
+	}
+
+	public function getOptionalInputShape(): array {
+		return [];
+	}
+
+	public function getOptionalInputShapeEnumValues(): array {
+		return [];
+	}
+
+	public function getOptionalInputShapeDefaults(): array {
+		return [];
+	}
+
+	public function getOutputShapeEnumValues(): array {
+		return [];
+	}
+
+	public function getOptionalOutputShape(): array {
+		return [];
+	}
+
+	public function getOptionalOutputShapeEnumValues(): array {
+		return [];
+	}
+
+	/**
+	 * @inheritDoc
+	 * @return array{sources: list<string>}
+	 * @throws \RuntimeException
+	 */
+	public function process(?string $userId, array $input, callable $reportProgress): array {
+		if ($userId === null) {
+			throw new \RuntimeException('User ID is required to process the prompt.');
+		}
+
+		if (!isset($input['prompt']) || !is_string($input['prompt'])) {
+			throw new \RuntimeException('Invalid input, expected "prompt" key with string value');
+		}
+
+		if (!isset($input['limit']) || !is_numeric($input['limit'])) {
+			throw new \RuntimeException('Invalid input, expected "limit" key with number value');
+		}
+		$limit = (int)$input['limit'];
+
+		if (
+			!isset($input['scopeType']) || !is_string($input['scopeType'])
+			|| !isset($input['scopeList']) || !is_array($input['scopeList'])
+			|| !isset($input['scopeListMeta']) || !is_string($input['scopeListMeta'])
+		) {
+			throw new \RuntimeException('Invalid input, expected "scopeType" key with string value, "scopeList" key with array value and "scopeListMeta" key with string value');
+		}
+
+		try {
+			ScopeType::validate($input['scopeType']);
+		} catch (\InvalidArgumentException $e) {
+			throw new \RuntimeException($e->getMessage(), intval($e->getCode()), $e);
+		}
+		if ($input['scopeType'] === ScopeType::SOURCE) {
+			throw new \InvalidArgumentException('Invalid scope type, source cannot be used to search');
+		}
+
+		// unscoped query
+		if ($input['scopeType'] === ScopeType::NONE) {
+			$response = $this->langRopeService->docSearch(
+				$userId,
+				$input['prompt'],
+				null,
+				null,
+				$limit,
+			);
+			if (isset($response['error'])) {
+				throw new \RuntimeException('No result in ContextChat response. ' . $response['error']);
+			}
+			return $this->processResponse($userId, $response);
+		}
+
+		// scoped query
+		$scopeList = array_unique($input['scopeList']);
+		if (count($scopeList) === 0) {
+			throw new \RuntimeException('Empty scope list provided, use unscoped query instead');
+		}
+
+		if ($input['scopeType'] === ScopeType::PROVIDER) {
+			/** @var array<string> $scopeList */
+			$processedScopes = $scopeList;
+			$this->logger->debug('No need to index sources, querying ContextChat', ['scopeType' => $input['scopeType'], 'scopeList' => $processedScopes]);
+		} else {
+			// this should never happen
+			throw new \InvalidArgumentException('Invalid scope type');
+		}
+
+		if (count($processedScopes) === 0) {
+			throw new \RuntimeException('No supported sources found in the scope list, extend the list or use unscoped query instead');
+		}
+
+		$response = $this->langRopeService->docSearch(
+			$userId,
+			$input['prompt'],
+			$input['scopeType'],
+			$processedScopes,
+			$limit,
+		);
+
+		return $this->processResponse($userId, $response);
+	}
+
+	/**
+	 * Validate and enrich sources JSON strings of the response
+	 *
+	 * @param string $userId
+	 * @param array $response
+	 * @return array{sources: list<string>}
+	 * @throws \RuntimeException
+	 */
+	private function processResponse(string $userId, array $response): array {
+		if (isset($response['error'])) {
+			throw new \RuntimeException('Error received in ContextChat document search request: ' . $response['error']);
+		}
+		if (!array_is_list($response)) {
+			throw new \RuntimeException('Invalid response from ContextChat, expected a list: ' . json_encode($response));
+		}
+
+		if (count($response) === 0) {
+			$this->logger->info('No sources found in the response', ['response' => $response]);
+			return [
+				'sources' => [],
+			];
+		}
+
+		$sources = $response;
+		$jsonSources = array_filter(array_map(
+			fn ($source) => json_encode($source),
+			$this->metadataService->getEnrichedSources(
+				$userId,
+				...array_map(
+					fn ($source) => $source['source_id'] ?? null,
+					$sources,
+				),
+			),
+		), fn ($json) => is_string($json));
+
+		if (count($jsonSources) === 0) {
+			$this->logger->warning('No sources could be enriched', ['sources' => $sources]);
+		} elseif (count($jsonSources) !== count($sources)) {
+			$this->logger->warning('Some sources could not be enriched', ['sources' => $sources, 'jsonSources' => $jsonSources]);
+		}
+
+		return [
+			'sources' => $jsonSources,
+		];
+	}
+}