Skip to content

Commit

Permalink
Fix agent bug, remove alias (#1661)
Browse files Browse the repository at this point in the history
  • Loading branch information
NolanTrem authored Dec 5, 2024
1 parent fb2a221 commit ab1d205
Show file tree
Hide file tree
Showing 12 changed files with 172 additions and 210 deletions.
2 changes: 1 addition & 1 deletion js/sdk/package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion js/sdk/package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "r2r-js",
"version": "0.4.4",
"version": "0.4.5",
"description": "",
"main": "dist/index.js",
"browser": "dist/index.browser.js",
Expand Down
75 changes: 0 additions & 75 deletions js/sdk/src/models.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -17,80 +17,11 @@ export interface RefreshTokenResponse {
};
}

export interface GenerationConfig {
model?: string;
temperature?: number;
topP?: number;
maxTokensToSample?: number;
stream?: boolean;
functions?: Array<Record<string, any>>;
tools?: Array<Record<string, any>>;
addGenerationKwargs?: Record<string, any>;
apiBase?: string;
responseFormat?: string;
}

export interface HybridSearchSettings {
fullTextWeight: number;
semanticWeight: number;
fullTextLimit: number;
rrfK: number;
}

export interface ChunkSearchSettings {
useVectorSearch?: boolean;
useHybridSearch?: boolean;
filters?: Record<string, any>;
searchLimit?: number;
offset?: number;
selectedCollectionIds?: string[];
indexMeasure: IndexMeasure;
includeScores?: boolean;
includeMetadatas?: boolean;
probes?: number;
efSearch?: number;
hybridSearchSettings?: HybridSearchSettings;
searchStrategy?: string;
}

export interface KGSearchSettings {
useKgSearch?: boolean;
filters?: Record<string, any>;
selectedCollectionIds?: string[];
graphragMapSystemPrompt?: string;
kgSearchType?: "local";
kgSearchLevel?: number | null;
generationConfig?: GenerationConfig;
maxCommunityDescriptionLength?: number;
maxLlmQueriesForGlobalSearch?: number;
localSearchLimits?: Record<string, number>;
}

export enum KGRunType {
ESTIMATE = "estimate",
RUN = "run",
}

export interface KGCreationSettings {
kgRelationshipsExtractionPrompt?: string;
kgEntityDescriptionPrompt?: string;
forceKgCreation?: boolean;
entityTypes?: string[];
relationTypes?: string[];
extractionsMergeCount?: number;
maxKnowledgeRelationships?: number;
maxDescriptionInputLength?: number;
generationConfig?: GenerationConfig;
}

export interface KGEnrichmentSettings {
forceKgEnrichment?: boolean;
communityReportsPrompt?: string;
maxSummaryInputLength?: number;
generationConfig?: GenerationConfig;
leidenParams?: Record<string, any>;
}

export interface KGEntityDeduplicationSettings {
kgEntityDeduplicationType?: KGEntityDeduplicationType;
}
Expand Down Expand Up @@ -121,12 +52,6 @@ export interface R2RDocumentChunksRequest {
documentId: string;
}

export enum IndexMeasure {
COSINE_DISTANCE = "cosine_distance",
L2_DISTANCE = "l2_distance",
MAX_INNER_PRODUCT = "max_inner_product",
}

export interface RawChunk {
text: string;
}
27 changes: 11 additions & 16 deletions js/sdk/src/r2rClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,8 @@ import {
TokenInfo,
Message,
RefreshTokenResponse,
ChunkSearchSettings,
KGSearchSettings,
KGRunType,
KGCreationSettings,
KGEnrichmentSettings,
KGEntityDeduplicationSettings,
GenerationConfig,
RawChunk,
} from "./models";

Expand Down Expand Up @@ -1615,7 +1610,7 @@ export class r2rClient extends BaseClient {
async createGraph(
collection_id?: string,
run_type?: KGRunType,
graph_creation_settings?: KGCreationSettings | Record<string, any>,
graph_creation_settings?: Record<string, any>,
): Promise<Record<string, any>> {
this._ensureAuthenticated();

Expand Down Expand Up @@ -1643,7 +1638,7 @@ export class r2rClient extends BaseClient {
async enrichGraph(
collection_id?: string,
run_type?: KGRunType,
graph_enrichment_settings?: KGEnrichmentSettings | Record<string, any>,
graph_enrichment_settings?: Record<string, any>,
): Promise<any> {
this._ensureAuthenticated();

Expand Down Expand Up @@ -1866,7 +1861,7 @@ export class r2rClient extends BaseClient {
@feature("searchDocuments")
async searchDocuments(
query: string,
vector_search_settings?: ChunkSearchSettings | Record<string, any>,
vector_search_settings?: Record<string, any>,
): Promise<any> {
this._ensureAuthenticated();
const json_data: Record<string, any> = {
Expand Down Expand Up @@ -1894,8 +1889,8 @@ export class r2rClient extends BaseClient {
@feature("search")
async search(
query: string,
vector_search_settings?: ChunkSearchSettings | Record<string, any>,
graph_search_settings?: KGSearchSettings | Record<string, any>,
vector_search_settings?: Record<string, any>,
graph_search_settings?: Record<string, any>,
): Promise<any> {
this._ensureAuthenticated();

Expand Down Expand Up @@ -1926,9 +1921,9 @@ export class r2rClient extends BaseClient {
@feature("rag")
async rag(
query: string,
vector_search_settings?: ChunkSearchSettings | Record<string, any>,
graph_search_settings?: KGSearchSettings | Record<string, any>,
rag_generation_config?: GenerationConfig | Record<string, any>,
vector_search_settings?: Record<string, any>,
graph_search_settings?: Record<string, any>,
rag_generation_config?: Record<string, any>,
task_prompt_override?: string,
include_title_if_available?: boolean,
): Promise<any | AsyncGenerator<string, void, unknown>> {
Expand Down Expand Up @@ -1986,9 +1981,9 @@ export class r2rClient extends BaseClient {
@feature("agent")
async agent(
messages: Message[],
rag_generation_config?: GenerationConfig | Record<string, any>,
vector_search_settings?: ChunkSearchSettings | Record<string, any>,
graph_search_settings?: KGSearchSettings | Record<string, any>,
rag_generation_config?: Record<string, any>,
vector_search_settings?: Record<string, any>,
graph_search_settings?: Record<string, any>,
task_prompt_override?: string,
include_title_if_available?: boolean,
conversation_id?: string,
Expand Down
54 changes: 46 additions & 8 deletions js/sdk/src/types.ts
Original file line number Diff line number Diff line change
@@ -1,9 +1,3 @@
import {
ChunkSearchSettings,
GenerationConfig,
HybridSearchSettings,
} from "./models";

export interface UnprocessedChunk {
id: string;
document_id?: string;
Expand Down Expand Up @@ -134,6 +128,13 @@ export interface GraphResponse {
updated_at: string;
}

// Index types
export enum IndexMeasure {
COSINE_DISTANCE = "cosine_distance",
L2_DISTANCE = "l2_distance",
MAX_INNER_PRODUCT = "max_inner_product",
}

// Ingestion types
export interface IngestionResponse {
message: string;
Expand Down Expand Up @@ -184,8 +185,41 @@ export interface RelationshipResponse {
}

// Retrieval types
export interface ChunkSearchSettings {
index_measure?: IndexMeasure;
probes?: number;
ef_search?: number;
enabled?: boolean;
}

export interface GenerationConfig {
model?: string;
temperature?: number;
top_p?: number;
max_tokens_to_sample?: number;
stream?: boolean;
functions?: Array<Record<string, any>>;
tools?: Array<Record<string, any>>;
add_generation_kwargs?: Record<string, any>;
api_base?: string;
response_format?: string;
}

export interface HybridSearchSettings {
full_text_weight?: number;
semantic_weight?: number;
full_text_limit?: number;
rrf_k?: number;
}

export interface GraphSearchSettings {
generation_config?: GenerationConfig;
graphrag_map_system?: string;
graphrag_reduce_system?: string;
max_community_description_length?: number;
max_llm_queries_for_global_search?: number;
limits?: Record<string, any>;
enabled?: boolean;
}

export interface SearchSettings {
Expand All @@ -196,7 +230,7 @@ export interface SearchSettings {
limit?: number;
offset?: number;
include_metadata?: boolean;
include_scores: boolean;
include_scores?: boolean;
search_strategy?: string;
hybrid_settings?: HybridSearchSettings;
chunk_settings?: ChunkSearchSettings;
Expand All @@ -212,7 +246,11 @@ export interface VectorSearchResult {
metadata?: Record<string, any>;
}

export type KGSearchResultType = "entity" | "relationship" | "community" | "global";
export type KGSearchResultType =
| "entity"
| "relationship"
| "community"
| "global";

export interface GraphSearchResult {
content: any;
Expand Down
Loading

0 comments on commit ab1d205

Please sign in to comment.