sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-12-18 18:47:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 23

Initial commit of a functional but not yet elegant prototype for this concept

Browse source
This commit is contained in:
sabaimran 2024-11-28 17:28:23 -08:00
parent 9368699b2c
commit d91935c880
15 changed files with 455 additions and 150 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
src/interface/web/app/common/chatFunctions.ts
View file

@ -1,3 +1,4 @@
import { AttachedFileText } from "../components/chatInputArea/chatInputArea";
import {
CodeContext,
Context,
@ -16,6 +17,12 @@ export interface MessageMetadata {
turnId: string;
}
export interface GeneratedAssetsData {
images: string[];
excalidrawDiagram: string;
files: AttachedFileText[];
}
export interface ResponseWithIntent {
intentType: string;
response: string;
@ -84,6 +91,8 @@ export function processMessageChunk(
if (!currentMessage || !chunk || !chunk.type) return { context, onlineContext, codeContext };
console.log(`chunk type: ${chunk.type}`);
if (chunk.type === "status") {
console.log(`status: ${chunk.data}`);
const statusMessage = chunk.data as string;
@ -98,6 +107,20 @@ export function processMessageChunk(
} else if (chunk.type === "metadata") {
const messageMetadata = chunk.data as MessageMetadata;
currentMessage.turnId = messageMetadata.turnId;
} else if (chunk.type === "generated_assets") {
const generatedAssets = chunk.data as GeneratedAssetsData;
if (generatedAssets.images) {
currentMessage.generatedImages = generatedAssets.images;
}
if (generatedAssets.excalidrawDiagram) {
currentMessage.generatedExcalidrawDiagram = generatedAssets.excalidrawDiagram;
}
if (generatedAssets.files) {
currentMessage.generatedFiles = generatedAssets.files;
}
} else if (chunk.type === "message") {
const chunkData = chunk.data;
// Here, handle if the response is a JSON response with an image, but the intentType is excalidraw

9
src/interface/web/app/components/chatHistory/chatHistory.tsx
View file

@ -54,6 +54,12 @@ function TrainOfThoughtComponent(props: TrainOfThoughtComponentProps) {
const lastIndex = props.trainOfThought.length - 1;
const [collapsed, setCollapsed] = useState(props.completed);
// useEffect(() => {
// if (props.completed) {
// setCollapsed(true);
// }
// }), [props.completed];
return (
<div
className={`${!collapsed ? styles.trainOfThought + " shadow-sm" : ""}`}
@ -410,6 +416,9 @@ export default function ChatHistory(props: ChatHistoryProps) {
"inferred-queries": message.inferredQueries || [],
},
conversationId: props.conversationId,
images: message.generatedImages,
queryFiles: message.generatedFiles,
excalidrawDiagram: message.generatedExcalidrawDiagram,
turnId: messageTurnId,
}}
conversationId={props.conversationId}

12
src/interface/web/app/components/chatMessage/chatMessage.tsx
View file

@ -163,6 +163,7 @@ export interface SingleChatMessage {
conversationId: string;
turnId?: string;
queryFiles?: AttachedFileText[];
excalidrawDiagram?: string;
}
export interface StreamMessage {
@ -180,6 +181,10 @@ export interface StreamMessage {
inferredQueries?: string[];
turnId?: string;
queryFiles?: AttachedFileText[];
excalidrawDiagram?: string;
generatedFiles?: AttachedFileText[];
generatedImages?: string[];
generatedExcalidrawDiagram?: string;
}
export interface ChatHistoryData {
@ -264,6 +269,9 @@ interface ChatMessageProps {
onDeleteMessage: (turnId?: string) => void;
conversationId: string;
turnId?: string;
generatedImage?: string;
excalidrawDiagram?: string;
generatedFiles?: AttachedFileText[];
}
interface TrainOfThoughtProps {
@ -394,6 +402,10 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
setExcalidrawData(props.chatMessage.message);
}
if (props.chatMessage.excalidrawDiagram) {
setExcalidrawData(props.chatMessage.excalidrawDiagram);
}
// Replace LaTeX delimiters with placeholders
message = message
.replace(/\\\(/g, "LEFTPAREN")

195
src/khoj/database/models/__init__.py
View file

@ -1,7 +1,9 @@
import logging
import os
import re
import uuid
from random import choice
from typing import Dict, List, Optional, Union
from django.contrib.auth.models import AbstractUser
from django.contrib.postgres.fields import ArrayField
@ -11,9 +13,109 @@ from django.db.models.signals import pre_save
from django.dispatch import receiver
from pgvector.django import VectorField
from phonenumber_field.modelfields import PhoneNumberField
from pydantic import BaseModel as PydanticBaseModel
from pydantic import Field
logger = logging.getLogger(__name__)
class BaseModel(models.Model):
# Pydantic models for type Chat Message validation
class Context(PydanticBaseModel):
compiled: str
file: str
class CodeContextFile(PydanticBaseModel):
filename: str
b64_data: str
class CodeContextResult(PydanticBaseModel):
success: bool
output_files: List[CodeContextFile]
std_out: str
std_err: str
code_runtime: int
class CodeContextData(PydanticBaseModel):
code: str
result: CodeContextResult
class WebPage(PydanticBaseModel):
link: str
query: Optional[str] = None
snippet: str
class AnswerBox(PydanticBaseModel):
link: str
snippet: str
title: str
snippetHighlighted: List[str]
class PeopleAlsoAsk(PydanticBaseModel):
link: str
question: str
snippet: str
title: str
class KnowledgeGraph(PydanticBaseModel):
attributes: Dict[str, str]
description: str
descriptionLink: str
descriptionSource: str
imageUrl: str
title: str
type: str
class OrganicContext(PydanticBaseModel):
snippet: str
title: str
link: str
class OnlineContext(PydanticBaseModel):
webpages: Optional[Union[WebPage, List[WebPage]]] = None
answerBox: Optional[AnswerBox] = None
peopleAlsoAsk: Optional[List[PeopleAlsoAsk]] = None
knowledgeGraph: Optional[KnowledgeGraph] = None
organicContext: Optional[List[OrganicContext]] = None
class Intent(PydanticBaseModel):
type: str
query: str
memory_type: str = Field(alias="memory-type")
inferred_queries: Optional[List[str]] = Field(default=None, alias="inferred-queries")
class TrainOfThought(PydanticBaseModel):
type: str
data: str
class ChatMessage(PydanticBaseModel):
message: str
trainOfThought: List[TrainOfThought] = []
context: List[Context] = []
onlineContext: Dict[str, OnlineContext] = {}
codeContext: Dict[str, CodeContextData] = {}
created: str
images: Optional[List[str]] = None
queryFiles: Optional[List[Dict]] = None
excalidrawDiagram: Optional[str] = None
by: str
turnId: Optional[str]
intent: Optional[Intent] = None
automationId: Optional[str] = None
class DbBaseModel(models.Model):
created_at = models.DateTimeField(auto_now_add=True)
updated_at = models.DateTimeField(auto_now=True)
@ -21,7 +123,7 @@ class BaseModel(models.Model):
abstract = True
class ClientApplication(BaseModel):
class ClientApplication(DbBaseModel):
name = models.CharField(max_length=200)
client_id = models.CharField(max_length=200)
client_secret = models.CharField(max_length=200)
@ -67,7 +169,7 @@ class KhojApiUser(models.Model):
accessed_at = models.DateTimeField(null=True, default=None)
class Subscription(BaseModel):
class Subscription(DbBaseModel):
class Type(models.TextChoices):
TRIAL = "trial"
STANDARD = "standard"
@ -79,13 +181,13 @@ class Subscription(BaseModel):
enabled_trial_at = models.DateTimeField(null=True, default=None, blank=True)
class OpenAIProcessorConversationConfig(BaseModel):
class OpenAIProcessorConversationConfig(DbBaseModel):
name = models.CharField(max_length=200)
api_key = models.CharField(max_length=200)
api_base_url = models.URLField(max_length=200, default=None, blank=True, null=True)
class ChatModelOptions(BaseModel):
class ChatModelOptions(DbBaseModel):
class ModelType(models.TextChoices):
OPENAI = "openai"
OFFLINE = "offline"
@ -103,12 +205,12 @@ class ChatModelOptions(BaseModel):
)
class VoiceModelOption(BaseModel):
class VoiceModelOption(DbBaseModel):
model_id = models.CharField(max_length=200)
name = models.CharField(max_length=200)
class Agent(BaseModel):
class Agent(DbBaseModel):
class StyleColorTypes(models.TextChoices):
BLUE = "blue"
GREEN = "green"
@ -208,7 +310,7 @@ class Agent(BaseModel):
super().save(*args, **kwargs)
class ProcessLock(BaseModel):
class ProcessLock(DbBaseModel):
class Operation(models.TextChoices):
INDEX_CONTENT = "index_content"
SCHEDULED_JOB = "scheduled_job"
@ -231,24 +333,24 @@ def verify_agent(sender, instance, **kwargs):
raise ValidationError(f"A private Agent with the name {instance.name} already exists.")
class NotionConfig(BaseModel):
class NotionConfig(DbBaseModel):
token = models.CharField(max_length=200)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
class GithubConfig(BaseModel):
class GithubConfig(DbBaseModel):
pat_token = models.CharField(max_length=200)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
class GithubRepoConfig(BaseModel):
class GithubRepoConfig(DbBaseModel):
name = models.CharField(max_length=200)
owner = models.CharField(max_length=200)
branch = models.CharField(max_length=200)
github_config = models.ForeignKey(GithubConfig, on_delete=models.CASCADE, related_name="githubrepoconfig")
class WebScraper(BaseModel):
class WebScraper(DbBaseModel):
class WebScraperType(models.TextChoices):
FIRECRAWL = "Firecrawl"
OLOSTEP = "Olostep"
@ -321,7 +423,7 @@ class WebScraper(BaseModel):
super().save(*args, **kwargs)
class ServerChatSettings(BaseModel):
class ServerChatSettings(DbBaseModel):
chat_default = models.ForeignKey(
ChatModelOptions, on_delete=models.CASCADE, default=None, null=True, blank=True, related_name="chat_default"
)
@ -333,35 +435,35 @@ class ServerChatSettings(BaseModel):
)
class LocalOrgConfig(BaseModel):
class LocalOrgConfig(DbBaseModel):
input_files = models.JSONField(default=list, null=True)
input_filter = models.JSONField(default=list, null=True)
index_heading_entries = models.BooleanField(default=False)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
class LocalMarkdownConfig(BaseModel):
class LocalMarkdownConfig(DbBaseModel):
input_files = models.JSONField(default=list, null=True)
input_filter = models.JSONField(default=list, null=True)
index_heading_entries = models.BooleanField(default=False)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
class LocalPdfConfig(BaseModel):
class LocalPdfConfig(DbBaseModel):
input_files = models.JSONField(default=list, null=True)
input_filter = models.JSONField(default=list, null=True)
index_heading_entries = models.BooleanField(default=False)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
class LocalPlaintextConfig(BaseModel):
class LocalPlaintextConfig(DbBaseModel):
input_files = models.JSONField(default=list, null=True)
input_filter = models.JSONField(default=list, null=True)
index_heading_entries = models.BooleanField(default=False)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
class SearchModelConfig(BaseModel):
class SearchModelConfig(DbBaseModel):
class ModelType(models.TextChoices):
TEXT = "text"
@ -393,7 +495,7 @@ class SearchModelConfig(BaseModel):
bi_encoder_confidence_threshold = models.FloatField(default=0.18)
class TextToImageModelConfig(BaseModel):
class TextToImageModelConfig(DbBaseModel):
class ModelType(models.TextChoices):
OPENAI = "openai"
STABILITYAI = "stability-ai"
@ -430,7 +532,7 @@ class TextToImageModelConfig(BaseModel):
super().save(*args, **kwargs)
class SpeechToTextModelOptions(BaseModel):
class SpeechToTextModelOptions(DbBaseModel):
class ModelType(models.TextChoices):
OPENAI = "openai"
OFFLINE = "offline"
@ -439,22 +541,22 @@ class SpeechToTextModelOptions(BaseModel):
model_type = models.CharField(max_length=200, choices=ModelType.choices, default=ModelType.OFFLINE)
class UserConversationConfig(BaseModel):
class UserConversationConfig(DbBaseModel):
user = models.OneToOneField(KhojUser, on_delete=models.CASCADE)
setting = models.ForeignKey(ChatModelOptions, on_delete=models.CASCADE, default=None, null=True, blank=True)
class UserVoiceModelConfig(BaseModel):
class UserVoiceModelConfig(DbBaseModel):
user = models.OneToOneField(KhojUser, on_delete=models.CASCADE)
setting = models.ForeignKey(VoiceModelOption, on_delete=models.CASCADE, default=None, null=True, blank=True)
class UserTextToImageModelConfig(BaseModel):
class UserTextToImageModelConfig(DbBaseModel):
user = models.OneToOneField(KhojUser, on_delete=models.CASCADE)
setting = models.ForeignKey(TextToImageModelConfig, on_delete=models.CASCADE)
class Conversation(BaseModel):
class Conversation(DbBaseModel):
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
conversation_log = models.JSONField(default=dict)
client = models.ForeignKey(ClientApplication, on_delete=models.CASCADE, default=None, null=True, blank=True)
@ -468,8 +570,39 @@ class Conversation(BaseModel):
file_filters = models.JSONField(default=list)
id = models.UUIDField(default=uuid.uuid4, editable=False, unique=True, primary_key=True, db_index=True)
def clean(self):
# Validate conversation_log structure
try:
messages = self.conversation_log.get("chat", [])
for msg in messages:
ChatMessage.model_validate(msg)
except Exception as e:
raise ValidationError(f"Invalid conversation_log format: {str(e)}")
class PublicConversation(BaseModel):
def save(self, *args, **kwargs):
self.clean()
super().save(*args, **kwargs)
@property
def messages(self) -> List[ChatMessage]:
"""Type-hinted accessor for conversation messages"""
validated_messages = []
for msg in self.conversation_log.get("chat", []):
try:
# Clean up inferred queries if they contain None
if msg.get("intent") and msg["intent"].get("inferred-queries"):
msg["intent"]["inferred-queries"] = [
q for q in msg["intent"]["inferred-queries"] if q is not None and isinstance(q, str)
]
msg["message"] = str(msg.get("message", ""))
validated_messages.append(ChatMessage.model_validate(msg))
except ValidationError as e:
logger.warning(f"Skipping invalid message in conversation: {e}")
continue
return validated_messages
class PublicConversation(DbBaseModel):
source_owner = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
conversation_log = models.JSONField(default=dict)
slug = models.CharField(max_length=200, default=None, null=True, blank=True)
@ -499,12 +632,12 @@ def verify_public_conversation(sender, instance, **kwargs):
instance.slug = slug
class ReflectiveQuestion(BaseModel):
class ReflectiveQuestion(DbBaseModel):
question = models.CharField(max_length=500)
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE, default=None, null=True, blank=True)
class Entry(BaseModel):
class Entry(DbBaseModel):
class EntryType(models.TextChoices):
IMAGE = "image"
PDF = "pdf"
@ -541,7 +674,7 @@ class Entry(BaseModel):
raise ValidationError("An Entry cannot be associated with both a user and an agent.")
class FileObject(BaseModel):
class FileObject(DbBaseModel):
# Same as Entry but raw will be a much larger string
file_name = models.CharField(max_length=400, default=None, null=True, blank=True)
raw_text = models.TextField()
@ -549,7 +682,7 @@ class FileObject(BaseModel):
agent = models.ForeignKey(Agent, on_delete=models.CASCADE, default=None, null=True, blank=True)
class EntryDates(BaseModel):
class EntryDates(DbBaseModel):
date = models.DateField()
entry = models.ForeignKey(Entry, on_delete=models.CASCADE, related_name="embeddings_dates")
@ -559,12 +692,12 @@ class EntryDates(BaseModel):
]
class UserRequests(BaseModel):
class UserRequests(DbBaseModel):
user = models.ForeignKey(KhojUser, on_delete=models.CASCADE)
slug = models.CharField(max_length=200)
class DataStore(BaseModel):
class DataStore(DbBaseModel):
key = models.CharField(max_length=200, unique=True)
value = models.JSONField(default=dict)
private = models.BooleanField(default=False)

10
src/khoj/processor/conversation/anthropic/anthropic_chat.py
View file

@ -2,7 +2,7 @@ import json
import logging
import re
from datetime import datetime, timedelta
from typing import Dict, Optional
from typing import Dict, List, Optional
from langchain.schema import ChatMessage
@ -158,6 +158,10 @@ def converse_anthropic(
query_images: Optional[list[str]] = None,
vision_available: bool = False,
query_files: str = None,
generated_images: Optional[list[str]] = None,
generated_files: List[str] = None,
generated_excalidraw_diagram: Optional[str] = None,
additional_context: Optional[str] = None,
tracer: dict = {},
):
"""
@ -218,6 +222,10 @@ def converse_anthropic(
vision_enabled=vision_available,
model_type=ChatModelOptions.ModelType.ANTHROPIC,
query_files=query_files,
generated_excalidraw_diagram=generated_excalidraw_diagram,
generated_files=generated_files,
generated_images=generated_images,
additional_program_context=additional_context,
)
messages, system_prompt = format_messages_for_anthropic(messages, system_prompt)

10
src/khoj/processor/conversation/google/gemini_chat.py
View file

@ -2,7 +2,7 @@ import json
import logging
import re
from datetime import datetime, timedelta
from typing import Dict, Optional
from typing import Dict, List, Optional
from langchain.schema import ChatMessage
@ -168,6 +168,10 @@ def converse_gemini(
query_images: Optional[list[str]] = None,
vision_available: bool = False,
query_files: str = None,
generated_images: Optional[list[str]] = None,
generated_files: List[str] = None,
generated_excalidraw_diagram: Optional[str] = None,
additional_context: List[str] = None,
tracer={},
):
"""
@ -229,6 +233,10 @@ def converse_gemini(
vision_enabled=vision_available,
model_type=ChatModelOptions.ModelType.GOOGLE,
query_files=query_files,
generated_excalidraw_diagram=generated_excalidraw_diagram,
generated_files=generated_files,
generated_images=generated_images,
additional_program_context=additional_context,
)
messages, system_prompt = format_messages_for_gemini(messages, system_prompt)

4
src/khoj/processor/conversation/offline/chat_model.py
View file

@ -162,6 +162,8 @@ def converse_offline(
user_name: str = None,
agent: Agent = None,
query_files: str = None,
generated_files: List[str] = None,
additional_context: List[str] = None,
tracer: dict = {},
) -> Union[ThreadedGenerator, Iterator[str]]:
"""
@ -229,6 +231,8 @@ def converse_offline(
tokenizer_name=tokenizer_name,
model_type=ChatModelOptions.ModelType.OFFLINE,
query_files=query_files,
generated_files=generated_files,
additional_program_context=additional_context,
)
logger.debug(f"Conversation Context for {model}: {messages_to_print(messages)}")

10
src/khoj/processor/conversation/openai/gpt.py
View file

@ -1,7 +1,7 @@
import json
import logging
from datetime import datetime, timedelta
from typing import Dict, Optional
from typing import Dict, List, Optional
from langchain.schema import ChatMessage
@ -157,6 +157,10 @@ def converse(
query_images: Optional[list[str]] = None,
vision_available: bool = False,
query_files: str = None,
generated_images: Optional[list[str]] = None,
generated_files: List[str] = None,
generated_excalidraw_diagram: Optional[str] = None,
additional_context: List[str] = None,
tracer: dict = {},
):
"""
@ -219,6 +223,10 @@ def converse(
vision_enabled=vision_available,
model_type=ChatModelOptions.ModelType.OPENAI,
query_files=query_files,
generated_excalidraw_diagram=generated_excalidraw_diagram,
generated_files=generated_files,
generated_images=generated_images,
additional_program_context=additional_context,
)
logger.debug(f"Conversation Context for GPT: {messages_to_print(messages)}")

21
src/khoj/processor/conversation/prompts.py
View file

@ -180,6 +180,20 @@ Improved Prompt:
""".strip()
)
generated_image_attachment = PromptTemplate.from_template(
f"""
Here is the image you generated based on my query. You can follow-up with a general response to my query. Limit to 1-2 sentences.
""".strip()
)
generated_diagram_attachment = PromptTemplate.from_template(
f"""
The AI has successfully created a diagram based on the user's query and handled the request. Good job!
AI can follow-up with a general response or summary. Limit to 1-2 sentences.
""".strip()
)
## Diagram Generation
## --
@ -1031,6 +1045,13 @@ A:
""".strip()
)
additional_program_context = PromptTemplate.from_template(
"""
Here's some additional context about what happened while I was executing this query:
{context}
""".strip()
)
personality_prompt_safety_expert_lax = PromptTemplate.from_template(
"""

64
src/khoj/processor/conversation/utils.py
View file

@ -155,6 +155,9 @@ def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="A
elif chat["by"] == "khoj" and ("text-to-image" in chat["intent"].get("type")):
chat_history += f"User: {chat['intent']['query']}\n"
chat_history += f"{agent_name}: [generated image redacted for space]\n"
elif chat["by"] == "khoj" and chat.get("images"):
chat_history += f"User: {chat['intent']['query']}\n"
chat_history += f"{agent_name}: [generated image redacted for space]\n"
elif chat["by"] == "khoj" and ("excalidraw" in chat["intent"].get("type")):
chat_history += f"User: {chat['intent']['query']}\n"
chat_history += f"{agent_name}: {chat['intent']['inferred-queries'][0]}\n"
@ -211,6 +214,7 @@ class ChatEvent(Enum):
END_LLM_RESPONSE = "end_llm_response"
MESSAGE = "message"
REFERENCES = "references"
GENERATED_ASSETS = "generated_assets"
STATUS = "status"
METADATA = "metadata"
USAGE = "usage"
@ -223,7 +227,6 @@ def message_to_log(
user_message_metadata={},
khoj_message_metadata={},
conversation_log=[],
train_of_thought=[],
):
"""Create json logs from messages, metadata for conversation log"""
default_khoj_message_metadata = {
@ -232,6 +235,10 @@ def message_to_log(
}
khoj_response_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
# Filter out any fields that are set to None
user_message_metadata = {k: v for k, v in user_message_metadata.items() if v is not None}
khoj_message_metadata = {k: v for k, v in khoj_message_metadata.items() if v is not None}
# Create json log from Human's message
human_log = merge_dicts({"message": user_message, "by": "you"}, user_message_metadata)
@ -259,6 +266,9 @@ def save_to_conversation_log(
automation_id: str = None,
query_images: List[str] = None,
raw_query_files: List[FileAttachment] = [],
generated_images: List[str] = [],
raw_generated_files: List[FileAttachment] = [],
generated_excalidraw_diagram: str = None,
train_of_thought: List[Any] = [],
tracer: Dict[str, Any] = {},
):
@ -281,9 +291,11 @@ def save_to_conversation_log(
"automationId": automation_id,
"trainOfThought": train_of_thought,
"turnId": turn_id,
"images": generated_images,
"queryFiles": [file.model_dump(mode="json") for file in raw_generated_files],
"excalidrawDiagram": str(generated_excalidraw_diagram),
},
conversation_log=meta_log.get("chat", []),
train_of_thought=train_of_thought,
)
ConversationAdapters.save_conversation(
user,
@ -307,7 +319,7 @@ Khoj: "{inferred_queries if ("text-to-image" in intent_type) else chat_response}
def construct_structured_message(
message: str, images: list[str], model_type: str, vision_enabled: bool, attached_file_context: str
message: str, images: list[str], model_type: str, vision_enabled: bool, attached_file_context: str = None
):
"""
Format messages into appropriate multimedia format for supported chat model types
@ -363,6 +375,10 @@ def generate_chatml_messages_with_context(
model_type="",
context_message="",
query_files: str = None,
generated_images: Optional[list[str]] = None,
generated_files: List[FileAttachment] = None,
generated_excalidraw_diagram: str = None,
additional_program_context: List[str] = [],
):
"""Generate chat messages with appropriate context from previous conversation to send to the chat model"""
# Set max prompt size from user config or based on pre-configured for model and machine specs
@ -382,6 +398,7 @@ def generate_chatml_messages_with_context(
message_attached_files = ""
chat_message = chat.get("message")
role = "user" if chat["by"] == "you" else "assistant"
if chat["by"] == "khoj" and "excalidraw" in chat["intent"].get("type", ""):
chat_message = chat["intent"].get("inferred-queries")[0]
@ -402,7 +419,7 @@ def generate_chatml_messages_with_context(
query_files_dict[file["name"]] = file["content"]
message_attached_files = gather_raw_query_files(query_files_dict)
chatml_messages.append(ChatMessage(content=message_attached_files, role="user"))
chatml_messages.append(ChatMessage(content=message_attached_files, role=role))
if not is_none_or_empty(chat.get("onlineContext")):
message_context += f"{prompts.online_search_conversation.format(online_results=chat.get('onlineContext'))}"
@ -411,9 +428,18 @@ def generate_chatml_messages_with_context(
reconstructed_context_message = ChatMessage(content=message_context, role="user")
chatml_messages.insert(0, reconstructed_context_message)
role = "user" if chat["by"] == "you" else "assistant"
if chat.get("images") and role == "assistant":
# Issue: the assistant role cannot accept an image as a message content, so send it in a separate user message.
file_attachment_message = construct_structured_message(
message=prompts.generated_image_attachment.format(),
images=chat.get("images"),
model_type=model_type,
vision_enabled=vision_enabled,
)
chatml_messages.append(ChatMessage(content=file_attachment_message, role="user"))
message_content = construct_structured_message(
chat_message, chat.get("images"), model_type, vision_enabled, attached_file_context=query_files
chat_message, chat.get("images") if role == "user" else [], model_type, vision_enabled
)
reconstructed_message = ChatMessage(content=message_content, role=role)
@ -423,6 +449,7 @@ def generate_chatml_messages_with_context(
break
messages = []
if not is_none_or_empty(user_message):
messages.append(
ChatMessage(
@ -435,6 +462,31 @@ def generate_chatml_messages_with_context(
if not is_none_or_empty(context_message):
messages.append(ChatMessage(content=context_message, role="user"))
if generated_images:
messages.append(
ChatMessage(
content=construct_structured_message(
prompts.generated_image_attachment.format(), generated_images, model_type, vision_enabled
),
role="user",
)
)
if generated_files:
message_attached_files = gather_raw_query_files({file.name: file.content for file in generated_files})
messages.append(ChatMessage(content=message_attached_files, role="assistant"))
if generated_excalidraw_diagram:
messages.append(ChatMessage(content=prompts.generated_diagram_attachment.format(), role="assistant"))
if additional_program_context:
messages.append(
ChatMessage(
content=prompts.additional_program_context.format(context="\n".join(additional_program_context)),
role="assistant",
)
)
if len(chatml_messages) > 0:
messages += chatml_messages

22
src/khoj/processor/image/generate.py
View file

@ -12,7 +12,7 @@ from khoj.database.models import Agent, KhojUser, TextToImageModelConfig
from khoj.routers.helpers import ChatEvent, generate_better_image_prompt
from khoj.routers.storage import upload_image
from khoj.utils import state
from khoj.utils.helpers import ImageIntentType, convert_image_to_webp, timer
from khoj.utils.helpers import convert_image_to_webp, timer
from khoj.utils.rawconfig import LocationData
logger = logging.getLogger(__name__)
@ -34,14 +34,13 @@ async def text_to_image(
status_code = 200
image = None
image_url = None
intent_type = ImageIntentType.TEXT_TO_IMAGE_V3
text_to_image_config = await ConversationAdapters.aget_user_text_to_image_model(user)
if not text_to_image_config:
# If the user has not configured a text to image model, return an unsupported on server error
status_code = 501
message = "Failed to generate image. Setup image generation on the server."
yield image_url or image, status_code, message, intent_type.value
yield image_url or image, status_code, message
return
text2image_model = text_to_image_config.model_name
@ -53,6 +52,9 @@ async def text_to_image(
elif chat["by"] == "khoj" and "text-to-image" in chat["intent"].get("type"):
chat_history += f"Q: Prompt: {chat['intent']['query']}\n"
chat_history += f"A: Improved Prompt: {chat['intent']['inferred-queries'][0]}\n"
elif chat["by"] == "khoj" and chat.get("images"):
chat_history += f"Q: {chat['intent']['query']}\n"
chat_history += f"A: Improved Prompt: {chat['intent']['inferred-queries'][0]}\n"
if send_status_func:
async for event in send_status_func("**Enhancing the Painting Prompt**"):
@ -92,31 +94,29 @@ async def text_to_image(
logger.error(f"Image Generation blocked by OpenAI: {e}")
status_code = e.status_code # type: ignore
message = f"Image generation blocked by OpenAI due to policy violation" # type: ignore
yield image_url or image, status_code, message, intent_type.value
yield image_url or image, status_code, message
return
else:
logger.error(f"Image Generation failed with {e}", exc_info=True)
message = f"Image generation failed using OpenAI" # type: ignore
status_code = e.status_code # type: ignore
yield image_url or image, status_code, message, intent_type.value
yield image_url or image, status_code, message
return
except requests.RequestException as e:
logger.error(f"Image Generation failed with {e}", exc_info=True)
message = f"Image generation using {text2image_model} via {text_to_image_config.model_type} failed due to a network error."
status_code = 502
yield image_url or image, status_code, message, intent_type.value
yield image_url or image, status_code, message
return
# Decide how to store the generated image
with timer("Upload image to S3", logger):
image_url = upload_image(webp_image_bytes, user.uuid)
if image_url:
intent_type = ImageIntentType.TEXT_TO_IMAGE2
else:
intent_type = ImageIntentType.TEXT_TO_IMAGE_V3
if not image_url:
image = base64.b64encode(webp_image_bytes).decode("utf-8")
yield image_url or image, status_code, image_prompt, intent_type.value
yield image_url or image, status_code, image_prompt
def generate_image_with_openai(

181
src/khoj/routers/api_chat.py
View file

@ -77,6 +77,7 @@ from khoj.utils.helpers import (
)
from khoj.utils.rawconfig import (
ChatRequestBody,
FileAttachment,
FileFilterRequest,
FilesFilterRequest,
LocationData,
@ -771,6 +772,11 @@ async def chat(
file_filters = conversation.file_filters if conversation and conversation.file_filters else []
attached_file_context = gather_raw_query_files(query_files)
generated_images: List[str] = []
generated_files: List[FileAttachment] = []
generated_excalidraw_diagram: str = None
additional_context_for_llm_response: List[str] = []
if conversation_commands == [ConversationCommand.Default] or is_automated_task:
chosen_io = await aget_data_sources_and_output_format(
q,
@ -876,21 +882,17 @@ async def chat(
async for result in send_llm_response(response, tracer.get("usage")):
yield result
await sync_to_async(save_to_conversation_log)(
q,
response_log,
user,
meta_log,
user_message_time,
intent_type="summarize",
client_application=request.user.client_app,
conversation_id=conversation_id,
query_images=uploaded_images,
train_of_thought=train_of_thought,
raw_query_files=raw_query_files,
tracer=tracer,
summarized_document = FileAttachment(
name="Summarized Document",
content=response_log,
type="text/plain",
size=len(response_log.encode("utf-8")),
)
return
async for result in send_event(ChatEvent.GENERATED_ASSETS, {"files": [summarized_document.model_dump()]}):
yield result
generated_files.append(summarized_document)
custom_filters = []
if conversation_commands == [ConversationCommand.Help]:
@ -1079,6 +1081,7 @@ async def chat(
async for result in send_event(ChatEvent.STATUS, f"**Ran code snippets**: {len(code_results)}"):
yield result
except ValueError as e:
additional_context_for_llm_response.append(f"Failed to run code")
logger.warning(
f"Failed to use code tool: {e}. Attempting to respond without code results",
exc_info=True,
@ -1116,51 +1119,36 @@ async def chat(
if isinstance(result, dict) and ChatEvent.STATUS in result:
yield result[ChatEvent.STATUS]
else:
generated_image, status_code, improved_image_prompt, intent_type = result
generated_image, status_code, improved_image_prompt = result
inferred_queries.append(improved_image_prompt)
if generated_image is None or status_code != 200:
content_obj = {
"content-type": "application/json",
"intentType": intent_type,
"detail": improved_image_prompt,
"image": None,
}
async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
additional_context_for_llm_response.append(f"Failed to generate image with {improved_image_prompt}")
async for result in send_event(ChatEvent.STATUS, f"Failed to generate image"):
yield result
return
else:
generated_images.append(generated_image)
# content_obj = {
# "intentType": intent_type,
# "inferredQueries": [improved_image_prompt],
# "image": generated_image,
# }
# async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
# yield result
# return
await sync_to_async(save_to_conversation_log)(
q,
generated_image,
user,
meta_log,
user_message_time,
intent_type=intent_type,
inferred_queries=[improved_image_prompt],
client_application=request.user.client_app,
conversation_id=conversation_id,
compiled_references=compiled_references,
online_results=online_results,
code_results=code_results,
query_images=uploaded_images,
train_of_thought=train_of_thought,
raw_query_files=raw_query_files,
tracer=tracer,
)
content_obj = {
"intentType": intent_type,
"inferredQueries": [improved_image_prompt],
"image": generated_image,
}
async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
yield result
return
async for result in send_event(
ChatEvent.GENERATED_ASSETS,
{
"images": [generated_image],
},
):
yield result
if ConversationCommand.Diagram in conversation_commands:
async for result in send_event(ChatEvent.STATUS, f"Creating diagram"):
yield result
intent_type = "excalidraw"
inferred_queries = []
diagram_description = ""
@ -1184,62 +1172,59 @@ async def chat(
if better_diagram_description_prompt and excalidraw_diagram_description:
inferred_queries.append(better_diagram_description_prompt)
diagram_description = excalidraw_diagram_description
generated_excalidraw_diagram = diagram_description
async for result in send_event(
ChatEvent.GENERATED_ASSETS,
{
"excalidrawDiagram": excalidraw_diagram_description,
},
):
yield result
else:
error_message = "Failed to generate diagram. Please try again later."
async for result in send_llm_response(error_message, tracer.get("usage")):
additional_context_for_llm_response.append(
f"AI attempted to programmatically generate a diagram but failed due to a program issue. Generally, it is able to do so, but encountered a system issue this time. AI can suggest text description or rendering of the diagram or user can try again with a simpler prompt."
)
async for result in send_event(ChatEvent.STATUS, error_message):
yield result
await sync_to_async(save_to_conversation_log)(
q,
error_message,
user,
meta_log,
user_message_time,
inferred_queries=[better_diagram_description_prompt],
client_application=request.user.client_app,
conversation_id=conversation_id,
compiled_references=compiled_references,
online_results=online_results,
code_results=code_results,
query_images=uploaded_images,
train_of_thought=train_of_thought,
raw_query_files=raw_query_files,
tracer=tracer,
)
return
# content_obj = {
# "intentType": intent_type,
# "inferredQueries": inferred_queries,
# "image": diagram_description,
# }
content_obj = {
"intentType": intent_type,
"inferredQueries": inferred_queries,
"image": diagram_description,
}
# await sync_to_async(save_to_conversation_log)(
# q,
# excalidraw_diagram_description,
# user,
# meta_log,
# user_message_time,
# intent_type="excalidraw",
# inferred_queries=[better_diagram_description_prompt],
# client_application=request.user.client_app,
# conversation_id=conversation_id,
# compiled_references=compiled_references,
# online_results=online_results,
# code_results=code_results,
# query_images=uploaded_images,
# train_of_thought=train_of_thought,
# raw_query_files=raw_query_files,
# generated_images=generated_images,
# tracer=tracer,
# )
await sync_to_async(save_to_conversation_log)(
q,
excalidraw_diagram_description,
user,
meta_log,
user_message_time,
intent_type="excalidraw",
inferred_queries=[better_diagram_description_prompt],
client_application=request.user.client_app,
conversation_id=conversation_id,
compiled_references=compiled_references,
online_results=online_results,
code_results=code_results,
query_images=uploaded_images,
train_of_thought=train_of_thought,
raw_query_files=raw_query_files,
tracer=tracer,
)
async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
yield result
return
# async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
# yield result
# return
## Generate Text Output
async for result in send_event(ChatEvent.STATUS, f"**Generating a well-informed response**"):
yield result
llm_response, chat_metadata = await agenerate_chat_response(
defiltered_query,
meta_log,
@ -1259,6 +1244,10 @@ async def chat(
train_of_thought,
attached_file_context,
raw_query_files,
generated_images,
generated_files,
generated_excalidraw_diagram,
additional_context_for_llm_response,
tracer,
)

42
src/khoj/routers/helpers.py
View file

@ -1184,6 +1184,10 @@ def generate_chat_response(
train_of_thought: List[Any] = [],
query_files: str = None,
raw_query_files: List[FileAttachment] = None,
generated_images: List[str] = None,
raw_generated_files: List[FileAttachment] = [],
generated_excalidraw_diagram: str = None,
additional_context: List[str] = [],
tracer: dict = {},
) -> Tuple[Union[ThreadedGenerator, Iterator[str]], Dict[str, str]]:
# Initialize Variables
@ -1207,6 +1211,9 @@ def generate_chat_response(
query_images=query_images,
train_of_thought=train_of_thought,
raw_query_files=raw_query_files,
generated_images=generated_images,
raw_generated_files=raw_generated_files,
generated_excalidraw_diagram=generated_excalidraw_diagram,
tracer=tracer,
)
@ -1242,6 +1249,7 @@ def generate_chat_response(
user_name=user_name,
agent=agent,
query_files=query_files,
generated_files=raw_generated_files,
tracer=tracer,
)
@ -1268,6 +1276,10 @@ def generate_chat_response(
agent=agent,
vision_available=vision_available,
query_files=query_files,
generated_files=raw_generated_files,
generated_images=generated_images,
generated_excalidraw_diagram=generated_excalidraw_diagram,
additional_context=additional_context,
tracer=tracer,
)
@ -1291,6 +1303,10 @@ def generate_chat_response(
agent=agent,
vision_available=vision_available,
query_files=query_files,
generated_files=raw_generated_files,
generated_images=generated_images,
generated_excalidraw_diagram=generated_excalidraw_diagram,
additional_context=additional_context,
tracer=tracer,
)
elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
@ -1313,6 +1329,10 @@ def generate_chat_response(
query_images=query_images,
vision_available=vision_available,
query_files=query_files,
generated_files=raw_generated_files,
generated_images=generated_images,
generated_excalidraw_diagram=generated_excalidraw_diagram,
additional_context=additional_context,
tracer=tracer,
)
@ -1784,6 +1804,9 @@ class MessageProcessor:
self.references = {}
self.usage = {}
self.raw_response = ""
self.generated_images = []
self.generated_files = []
self.generated_excalidraw_diagrams = []
def convert_message_chunk_to_json(self, raw_chunk: str) -> Dict[str, Any]:
if raw_chunk.startswith("{") and raw_chunk.endswith("}"):
@ -1822,6 +1845,16 @@ class MessageProcessor:
self.raw_response += chunk_data
else:
self.raw_response += chunk_data
elif chunk_type == ChatEvent.GENERATED_ASSETS:
chunk_data = chunk["data"]
if isinstance(chunk_data, dict):
for key in chunk_data:
if key == "images":
self.generated_images = chunk_data[key]
elif key == "files":
self.generated_files = chunk_data[key]
elif key == "excalidraw_diagrams":
self.generated_excalidraw_diagrams = chunk_data[key]
def handle_json_response(self, json_data: Dict[str, str]) -> str | Dict[str, str]:
if "image" in json_data or "details" in json_data:
@ -1852,7 +1885,14 @@ async def read_chat_stream(response_iterator: AsyncGenerator[str, None]) -> Dict
if buffer:
processor.process_message_chunk(buffer)
return {"response": processor.raw_response, "references": processor.references, "usage": processor.usage}
return {
"response": processor.raw_response,
"references": processor.references,
"usage": processor.usage,
"images": processor.generated_images,
"files": processor.generated_files,
"excalidraw_diagrams": processor.generated_excalidraw_diagrams,
}
def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False):

1
tests/test_offline_chat_actors.py
View file

@ -22,7 +22,6 @@ from khoj.processor.conversation.offline.chat_model import (
filter_questions,
)
from khoj.processor.conversation.offline.utils import download_model
from khoj.processor.conversation.utils import message_to_log
from khoj.utils.constants import default_offline_chat_models

1
tests/test_online_chat_director.py
View file

@ -6,7 +6,6 @@ from freezegun import freeze_time
from khoj.database.models import Agent, Entry, KhojUser
from khoj.processor.conversation import prompts
from khoj.processor.conversation.utils import message_to_log
from tests.helpers import ConversationFactory, generate_chat_history, get_chat_api_key
# Initialize variables for tests